How to implement nested for loop in JAX for high-dimensional PDE solvers

[ToshiyukiBandai]

Hi all,

I would like to ask a question similar to #9191, which was not fully answered. I am writing a PDE solver in JAX, and it works fine but is too slow because of nested loops in it. My problem can be simplified into the code snippet below. Could anyone help me vectorize the nested for loops? I want to keep the matrix K_mat as a matrix, not a vector, if possible.

nx = 3
ny = 3

J = np.zeros(nx*ny)
K_mat = np.reshape(np.linspace(0, 1, nx*ny), (nx, ny))
alpha = np.linspace(0, 1, nx*ny)

for i in range(nx):
    for j in range(ny):
        k = j*nx + i
        value = K_mat[i, j] * alpha[k]
        J = J.at[k].set(value) 

[jakevdp]

One way to do this is via broadcasted indices; for example:

i = jnp.arange(nx)[None, :]  # shape=(1, nx)
j = jnp.arange(ny)[:, None]  # shape=(ny, 1)
k = j*nx + i
J = jnp.zeros(nx * ny).at[k].set(K_mat[i, j] * alpha[k])

[ToshiyukiBandai]

Hi jakevdp,

Thank you so much for your quick response! It looks great. I will try this with my PDE solver.

EDIT: It worked pretty well.

[patrick-kidger]

Jake's already answered this, but you may also find this nonlinear heat equation example a useful reference point.

[ToshiyukiBandai]

Hi Patrick,

Thank you for your comment. I read your thesis and it was quite informative!