jax.lax.scan using lists of arrays

[pvasired]

Is it possible to use jax.lax.scan to rewrite a for loop across several lists of jax.numpy arrays, where all lists have the same length but the elements of the lists may have different shapes?

For example:

def example_fun(arg):
    val = 0
    def body_fun(carry, x):
        a, b = x
        carry = carry + jnp.sum(a) + jnp.sum(b)

        return carry, ()

    val, _ = jax.lax.scan(body_fun, val, arg)

    return val 

arrays1 = [jnp.array([1, 2, 3]), jnp.array([4, 5, 6])]
arrays2 = [jnp.array([7, 8]), jnp.array([9, 10])]

example_fun((arrays1, arrays2))

What I would want this toy example to do is return 0 + (sum(1, 2, 3) + sum(7, 8)) + (sum(4, 5, 6) + sum(9, 10)) = 21 + 34 = 55.

This seems like it is possible since lists are pytrees (correct?), but I keep getting error messages about mismatched first axes dimensions when I try to do this. I don't understand where this error is coming from because both lists arrays1 and arrays2 have the same length (2). Specifically, the error is:

ValueError: scan got values with different leading axis sizes: 3, 3, 2, 2.

I need my inputs to be lists or something that can handle jagged shapes because in my real work, the elements of the lists do not have consistent shape. I cannot zero pad the arrays to make them consistent shape because that makes the code break (makes some matrices singular).

I appreciate any help I can get here! This is the first time I am using jax.lax.scan so I may be just doing something totally wrong.

[jakevdp]

scan is able to scan over arrays, not over lists. So one way you could accomplish what you want is to stack your inputs into arrays so that you're scanning over the first axis of each:

arrays1 = jnp.stack([jnp.array([1, 2, 3]), jnp.array([4, 5, 6])])
arrays2 = jnp.stack([jnp.array([7, 8]), jnp.array([9, 10])])

example_fun((arrays1, arrays2))
# 55

[jakevdp]

I believe @mattjj and @apaszke were experimenting with ragged arrays recently... tagging them here for visibility. Thanks for the interesting question! These sorts of out-of-the-box applications of JAX are really helpful in pushing development & improvement of the package.

[smartalecH]

Hi @jakevdp -- just to confirm, ragged arrays still aren't supported, correct? Are there other "tricks" to scan over a container which may hold arrays of different shapes (other than padding or python for-loops)? Thanks!

[jakevdp]

JAX still doesn't support ragged arrays: your best options are either to loop over a list of arrays, or to create a padded representation of the ragged array in order to use constructs like scan or vmap.

[MRiabov]

Hi @jakevdp -- just to confirm, ragged arrays still aren't supported, correct? Are there other "tricks" to scan over a container which may hold arrays of different shapes (other than padding or python for-loops)? Thanks!

@smartalecH, as somebody who tries to do it too, there is a neat solution: if you have arrays with different shapes except the first one, say:
[(10,6,6),
(250,6,6),
1000,6,6)]
You could stack along the first dimension, then vmap/scan and then slice the arrays in places of stack to get your variables back.
Would be neat if Jax supported this out of the box, somehow.

[jakevdp]

This works, so long as semantically what you want is a vmap of a concat. It doesn't work in general (for example if you are running an aggregate like a mean within each batch).

As for "out of the box" support: if these are the semantics you want, you can do this in a few lines with concat followed by vmap followed by split.