how to change the "at" operation to speed up jax

[ywsslr]

I have a loop to change big matrics' value,but only using "at" to change that.That's slower more than runing by numpy(I used gpu).So I want to konw how to chage the operation to speed up its runing.Belows is rest of my function.

seis = jnp.zeros((nt,len(sensors[0])))
p_t = jnp.zeros(v.shape)
p_t_1 = jnp.zeros(v.shape)
for t in range(nt):
    s[isz][isx] = signal[t]
    p_t = temp1*p_t_1 - temp2*p_t - beta_dt*s
    for i in [-4, -3, -2, -1, 1, 2, 3, 4]:
        p_t = p_t.at[4:l-4,4:l-4].add(alpha[4:l-4,4:l-4] * c[int(np.fabs(i))] * (p_t_1[4:l-4,4+i:l-4+i]+p_t_1[4+i:l-4+i,4:l-4]))
    seis = seis.at[t].set(p_t[irz,irx])
    p_t, p_t_1 = p_t_1, p_t

[jakevdp]

The best approach to speeding this up will involve rewriting it to avoid for loops. JAX (like NumPy) gets best performance from explicitly vectorized operations; as a simple example, this will be slow:

for i in range(x.shape[0]):
  x = x.at[i].set(2 * i)

While this will be fast:

i = jnp.arange(len(x))
x = x.at[i].set(2 * i)

I would suggest trying to use this sort of vectorized indexing to express your algorithm.