What is the best way to parallelize complex nested loops

[nkupelioglu]

Hello everyone,
I have got a nested loop that is indexed in a very complex way with slices and skips.

acc_idx = 0
for i in range(w):
    for j in range(h):
        for k in range(o):
            for n in range(b):
                matrix[start_idx + (i * h) + acc_idx + w, n, k, :, :] += x[n, k, j:j+h, i:i+w]
        acc_idx += w * h 

where matrix is of rank 5 and x is of rank 4, which are also fairly large. The reason I am pointing out this is that, I would like to implement them in an inplace way. (i.e. lax.dynamic_slice(?) )
Every update on this nested-loop is independent between each iteration of each loop, so technically it should be possible to unroll everything here. I am fairly new to Jax, so I wanted to ask how you would implement this (i.e. whether with nested vmaps, fori_loop), for least memory consumption and best timing.

[nestordemeure]

Usually that kind of access pattern (start_idx + (i * h) + acc_idx + w) is a sign of your matrix having some flattened dimensions.

My usual approach for those ugly loops is to reshape the data (matrix and x) to simplify the indexing, take the slice I need then either apply several layer of vmap (slightly error prone) or one xmap (much better as I can use a single function and give names to the sizes).

You cannot modify a jax array in place but there are two workaround:

• if you data starts in a numpy array, you can do something like matrix[:] = kernel(matrix,x)
• if your data is in a jax array, you can use the donate_buffer argument of the jit function when jitting your kernel (it will return an array but reuse the memory)