Making good use of fori_loop when shapes change

[fuglede]

Following up on @jakevdp's comment here, I have a case where I'm wondering of fori_loop will apply or not.

Concretely, I have something like

def f(X, k):
    state1 = jnp.empty((0, 0))
    state2 = X
    for _ in range(k):
        state1, state2 = update_state(state1, state2)
    return state1, state2

f_jit = jax.jit(f, static_argnames=("k",))

where X is an ndim 2 array, with shape about (1000, 1000), and k is about 200. More importantly, the method update_state is a JITable function which picks a column from state2 with a given property, removes it from state2 and uses it extend to state1 (and hopefully the details won't be too important). In particular, after one application of update_state, the shape of state1 changes from (n, m) to (n + 1, m + 1) with its original "top-left" (n, m)-block unchanged, and the shape of state2 changes from (n, m) to (n, m - 1), with one column being removed; which column to remove can not be determined a priori.

And in particular, this means that update_state does not conserve shapes, so jax.lax.fori_loop really doesn't like it, and therefore my question becomes: Is this an instance of some general pattern where I can use JAX intrinsics a la fori_loop or something else to reduce compile times?

[jakevdp]

In general, I think the best solution here would be to change update_state so that it returns padded arrays of the same shape as the input, and then use a fori_loop. Without more information about what update_state does, though, it's hard to say how feasible that is.

In other words, I think this is a case where the details are important.