Jax array value assignment takes very long runtime

[felix0901]

Hi All,
I found that in my code, the value assignment takes a very long runtime (one line of code takes >95% of the total runtime). I am executing value assignments on a very large array. Also, I am aware of the value assignment will create a new array which might be the cause of this issue.
I am currently doing it this way array = array.at[idxs].set(new_value)

I am wondering if you could give me some suggestions on how could I relieve the overhead of this?
Are there some tricks or walk-around I could use to reach the same goal with a much lighter implementation?

[jakevdp]

The recommended approach is to wrap your updates in jax.jit. From the jnp.ndarray.at documentation:

None of the x.at expressions modify the original x; instead they return a modified copy of x. However, inside a jit compiled function, expressions like x = x.at[idx].set(y) are guaranteed to be applied in-place.