How to perform a strided dynamic update slice?

[EdwardRaff]

To deal with sliding window functions, I've created some code that batches the sliding window into big chunks. So a (3,3) kernel would operate over 3*3=9 total views of the original matrix. This way I don't create a huge temporary array of all sliding windows at once (causing OOM due to intermediate representation size in JAX), and I don't get overhead/blowup of doing a scan one window at a time.

The only part that is missing is combining the stack of (9, C_out, H//3, W//3) tensors back into one uber tensor (C_out, H, W) where the values have been put in their correct place. I can do this with a strided update like this:

C_out =2
H, W = 11, 11
Q = jnp.zeros((C_out, H*3,W*3))
S = jnp.arange((9* C_out*H*W)).reshape(9, C_out, H, W)
for z in range(9):
    i, j = jnp.mod(z,kh), jnp.mod(z//kw,kh)
    Q = Q.at[:, i::3,j::3].add(S[z,:,:,:])

However, this isn't going to run nicely without a jit. The problem is hitting the Q.at[:, i::3,j::3] portion with a dynamic slice. I can write the loop for this as

def bod(ij, val):
    i, j = jnp.mod(z,kh), jnp.mod(z//kw,kh) #[0, kw*kh] -> tuples of module arithmatic (0,0), (1, 0), ...
    val = val.at[:, i::3,j::3].add(S[ij,:,:,:])
    return val
jax.lax.fori_loop(0, kw*kh, bod, Q)

But I get an error about static static start/stop/step. The jax.lax.dynamic_update_slice seems like the correct function, but the documentation has no indication on how to make this perform a strided dynamic update.

Any alternative approaches to the interlevade reconstruction of the 9 tensors is also appreciated. I need to interleave the rows and columns, e.g., something like the below index pattern:

Array([[0, 1, 2, 0, 1, 2, 0, 1, 2],
       [3, 4, 5, 3, 4, 5, 3, 4, 5],
       [6, 7, 8, 6, 7, 8, 6, 7, 8],
       [0, 1, 2, 0, 1, 2, 0, 1, 2],
       [3, 4, 5, 3, 4, 5, 3, 4, 5],
       [6, 7, 8, 6, 7, 8, 6, 7, 8],
       [0, 1, 2, 0, 1, 2, 0, 1, 2],
       [3, 4, 5, 3, 4, 5, 3, 4, 5],
       [6, 7, 8, 6, 7, 8, 6, 7, 8]], dtype=int32)

[dfm]

I'm not sure I totally understand the question or issues that you're encountering, but something like your first example works just fine under jit:

@jax.jit
def fun(S, Q):
  for z in range(9):
    i, j = np.mod(z,3), np.mod(z//3,3)  # <- Using 'np' not 'jnp' here
    Q = Q.at[:, i::3,j::3].add(S[z,:,:,:])
  return Q

C_out = 2
H, W = 11, 11
Q = jnp.zeros((C_out, H*3,W*3))
S = jnp.arange((9* C_out*H*W)).reshape(9, C_out, H, W)
Q = fun(S, Q)

Depending on your actual use case, you might also be able to use scan with the unroll parameter set >1, or (even better!) perhaps you can rewrite your problem as a convolution and use JAX's convolution implementation!

[EdwardRaff]

Yea doesn't seem like JAX can handle the for z in range(9) being converted into a fori or a scan, even with unroll=True, i,j become dynamic and I get the same error:

Array slice indices must have static start/stop/step to be used with NumPy indexing syntax. Found slice(Traced<ShapedArray(int32[])>with<DynamicJaxprTrace(level=1/0)>, None, 3). To index a statically sized array at a dynamic position, try lax.dynamic_slice/dynamic_update_slice (JAX does not support dynamically sized arrays within JIT compiled functions).

[dfm]

Can you share a minimal reproduction showing the issue you're having now (i.e. when the suggestion didn't fix your code)? Given what you've shared so far, that should work (even if it's not necessarily the most performant).

I always get tripped up by when I should use np instead of jnp!

My rule of thumb is that you want to use np when interacting with static quantities that are known at tracing time (i.e. shapes, dtypes, and in this case, indices). In this case, you should be able to print(i, j) and see ints not Traced objects.

[EdwardRaff]

C_out =2
H, W = 11, 11
Q = jnp.zeros((C_out, H*3,W*3))
S = jnp.arange((9* C_out*H*W)).reshape(9, C_out, H, W)

def minimal_func(kernel_size):
    kh, kw = kernel_size

    def f(carry, ij):
        i, j = np.mod(ij,kh), np.mod(ij//kw,kh)
        carry = carry.at[:, i::kh,j::kh].set(S[ij,:,:,:])
        return carry, jnp.zeros(1) #last arg not used
    return jax.lax.scan(f, Q, jnp.arange(kh*kw))
minimal_func((3,3)) #explodes here 

The minimum example above is! Unfortunately, once this is put into a scan, JAX thinks ij is a Trace instead of a statically known value.

[dfm]

Right, but you said it didn't work in a for loop. What was the problem there? It's not obvious to me why you want to use a scan here. I suggested it as something you can reach for in some related cases depending on your specific use case, but the easiest approach is to use a for loop here.

[EdwardRaff]

Sorry if I misstated. The for loop compiles and runs properly. However, for loops in JAX often compile poorly (e.g., if I wanted to do a (15,15) kernel shape). So ideally I could use a scan to make this maximally efficient and because it would allow me to merge the construction of the output with the calculation of the intermediate results (not included in minimum example b/c they are necessary to the issue for compilation).