How to avoid loops to make paralell filing of a 2D array

[jecampagne]

Hello,
here is a very simplified version of my use-case:

nx, ny= 20,30
x_val = np.linspace(-np.pi, np.pi,nx)
y_val = np.linspace(-np.pi/2,np.pi/2,ny)
result1 = np.zeros(shape=(nx,ny))
result2 = np.zeros(shape=(nx,ny))

for i,x in enumerate(x_val):
  for j,y in enumerate(y_val):
  # perform computations leading to 
   result1[i,j] = np.sin(x)*np.cos(y)
   result2[i,j] = np.sin(x*2)*np.cos(y**2)

# get result1 & result2

This code seems horrible and I only use np.sin/np.cos fucntions to get something simple but the code is fairly complicated to get the values of result1[i,j] and result2[i,j].

I would like to now how I can use jax.lax.scan or for_loop to get this use-case using the powerful jax/lax paralell computing. Any help is welcome.

[jakevdp]

It's hard to give a definite answer without more knowledge of what ... represents, but generally in JAX (as in numpy) you can use broadcasting for these kinds of operations. For example, if ... represents multiplication, you could do it like this:

import numpy as np

x_val = np.linspace(0, 100, 10)
y_val = np.linspace(0, 100, 20)

result = x_val[:, None] * y_val[None, :]
print(result.shape)
# (10, 20)

[jecampagne]

Sorry, the "# perform computations" is rather complicated computation based on "x" and "y" that are obtained by the 2 nested loops, and i,j is there to access the results1 & results2 buffer. So, this is not just what your code would give.

[jakevdp]

Makes sense - if you want help with your specific problem, I would suggest putting together a minimal reproducible example of what you're trying to do (perhaps replacing the full computation with something simpler that's of the same form). Otherwise, as we've seen, it's difficult to guess at your intent.

[YouJiacheng]

I think that is exactly what you need.
Just use jax.vmap, no need for jax.lax.scan.

def f(x, y):
    return jnp.sin(x) * jnp.cos(y) # any complicated computation based on a pair of elements

y_vmap_f = jax.vmap(f, (None, 0))
xy_vmap_f = jax.vmap(y_vmap_f, (0, None))

nx, ny = 20, 10
xs = jnp.arange(nx)
ys = jnp.arange(ny)

out = xy_vmap_f(xs, ys)

# verification of correctness
assert out.shape == (nx, ny)
for i in range(nx):
    for j in range(ny):
        assert out[i, j] == f(xs[i], ys[j])