Best way to implement JIT'able nested for-loop, where length of inner loop is the value of the outer loop

[rdaems]

Hi all,

What would be the best way to implement a JIT'able version of this example, where n can be static.

@partial(jax.jit, static_argnames=('n',))
def f(n, a):
    xs = []
    for i in range(n):
        x = 0.
        for j in range(i):
            x += np.exp(- (i - j)) * a[j]
        xs.append(x)
    xs = jnp.array(xs)
    return xs

This works, but takes a very long time to trace (for large n).

[jakevdp]

If you want to avoid long compile times, you should avoid using Python for loops. One strategy is to replace them with lax.fori_loop, but the length of such loops must be static (so the size of the inner loop could not depend on the counter of the outer loop).

In this case, I'd probably avoid that altogether and instead express your computation in a vectorized manner – then it will be both jit-compatible and performant. For example:

@partial(jax.jit, static_argnames=['n'])
def f(n, a):
  i = jnp.arange(n)[:, None]
  j = jnp.arange(n)[None, :]
  x = jnp.zeros((n, n))
  x = x.at[i, j].set(jnp.exp(-(i - j)) * a[j])
  return jnp.where(i <= j, 0, x).sum(1)

[rdaems]

Thanks for your suggestion.
We only need the upper triangular of x so this doubles the memory cost, right?

I managed to calculate all values in the upper triangular with a vmap with something like this:

def element(i, j):
    return jnp.exp(- (i - j)) * a[j]
ii = jnp.array([i for j in range(n) for i in range(j)], dtype=int)
jj = jnp.array([i for i in range(n) for _ in range(i)], dtype=int)
elements = jax.vmap(element)(ii, jj)

xs = jnp.array([elements[i * (i - 1) // 2:(i + 1) * i // 2].sum() for n in range(0, n)])

That's a bit faster, but the final sum still makes tracing slow.
Maybe I can try lax.fori_loop with dynamic indexing for that?

And of course this is more complicated than I would like, and less elegant than your solution.

[rdaems]

I guess there is no way around constructing the full x matrix? Since "all JAX arrays inside JIT must have statically known size".

[jakevdp]

This doubles the memory cost, yes, but that is a small price to pay in order to use the vector core of your accelearator. Any loopy implementation will be very slow if you're using GPU or TPU, and the only alternative is doing the loops host-side (as in both your solutions), which will lead to long compilation times.

If you're only targeting CPU, then it may be worth worrying about, and you could maybe do better by re-expressing this in terms of triu_indices and segment_sum, but even then I'm not sure the added complexity would be a net benefit performance-wise.

[rdaems]

I see. Thanks a lot for your help!