scan over arrays with different shapes

[fishjojo]

I'd like to perform a bunch of einsums for arrays with different shapes. I tried to use scan, but it seems not working. The script is as follows

import jax
from jax import numpy as jnp

aa = [jnp.zeros((5,5,2)), jnp.zeros((5,5,3))]
bb = [jnp.zeros((2,)), jnp.zeros((3,))]

def body(carry, xs):
    a, b = xs
    out = jnp.einsum('ijk,k->ij', a, b)
    return carry, out

carry, out = jax.lax.scan(body, None, (aa, bb))

and I got the error ValueError: scan got values with different leading axis sizes: 5, 5, 2, 3.

Am I doing it wrong or is there a workaround?
The loop is quite large in the real application so JIT won't work due to long compiling time.
I'd also like the function to be both forward and reverse differentiable.

Thank you in advance.

[jakevdp]

Hi - thanks for the question. scan can only be used to scan over JAX arrays, not Python lists. Since JAX does not have support for ragged arrays, it means that your use-case cannot be implemented via scan.

For what you have in mind, I'd suggest a simple Python loop. for example:

import jax
from jax import numpy as jnp

aa = [jnp.zeros((5,5,2)), jnp.zeros((5,5,3))]
bb = [jnp.zeros((2,)), jnp.zeros((3,))]

result = [jnp.einsum('ijk,k->ij', a, b) for a, b in zip(aa, bb)]

This will be compatible with both forward and reverse autodiff. If your lists are large, then wrapping the outer function in jit will likely lead to long compile times, because the for loop is flattened into a long sequence of individual operations. You can address this by not wrapping the outer function in jit.

Alternatively, you may be able to pad your values so that they can all be stored in a single array, then you could use scan and an outer jit without issues.

Hope that helps!

[fishjojo]

Thank you for the clarification. This is very helpful.