I have always same problem with error code 102 try both betas and dev version on macos (my iphone is always on newest beta also not working).

I have few loops that are part of my Deep Learning pipeline i got everything in MPSGraph but graph.for loop with scatter is extremely slow for my use case and there is no sane way to vectorize it. What is best way to update value in existing tensor/or to help Metal optimizer to vectorize this ops, maybe something like jax.at[idx].set? (Right now im using graph.for loop with scatter/gather ops)