Updating the element of the DeviceArray

[yiminghwang]

Hi, there,

I am going to update the element of the DeviceArray, but it runs much slower than Ndarry of NumPy, is there any way to improve the performance?

m,n = 100,100

def func1():
    A = jnp.zeros((m,n))
    start = time.time()
    for i in range(m):
        for j in range(n):
            A.at[i,j].set(i*j)
    end = time.time()
    return end-start

def func2():
    B = np.zeros((m,n))
    start = time.time()
    for i in range(m):
        for j in range(n):
            B[i,j]=i*j
    end = time.time()
    return end-start

run-time of func1:  23.88292932510376
run-time of func2:  0.001775979995727539

The above is just a simple test.
I aim to store the processed results in a DeviceArray.
So, should I use fori_loop + .at[indx].set instead of naive loop + at[indx].set?

[soraros]

Maybe you could share more information about your real™ use case?

There are some (hopefully useful) general suggestions I can give now:

1. .at[...].set(...) is functional, meaning you should write A = A.at[i, j].set(i * j)
2. One should prefer approaches without explicit loop at all, such as

@partial(jit, static_argnums=(0,))
def f(shape: tuple[int, int]) -> tuple[Array, Array]:
  I, J = jnp.indices(shape)
  A = jnp.zeros(shape).at[I, J].set(I * J)
  return A

or even simpler for your toy example

A = jnp.fromfunction(jnp.multiply, (m, n))

3. Jit compile your computation, and benchmark it carefully following this