Deal with static but ragged arrays in JAX

[renecotyfanboy]

Hi, this is a fairly specific question about a use case that I can't manage on my own.

The idea is simple: I want to work with a list of arrays of different sizes (corresponding to tabulated fluxes associated with a list of temperatures). The function interpolates these tables in two stages: I select 2 vectors around the desired temperature, I interpolate these two arrays to obtain, reducing them to two values, and I combine these two values.

The only issue is that each of this vector has a different length (due to inhomogeneous sampling), but they are static, and no intermediate vector is produced with a shape which is unknown at runtime. I have tried to pad everything to the same size, but this is not achievable in this specific use case, and try to think of way around to use vmap/scan or equivalent to get a fast evaluation.

Any thoughts on approaching this ?

[jakevdp]

Hi - I think there are two general possibilities here. First, you could store your ragged arrays in a Python list, where each element is a JAX array of different size. You wouldn't be able to vmap or scan over the list, but you could iterate using e.g. Python list comprehensions and XLA will generate efficient code for the resulting array operations. The only downside here is that these Python loops will be unrolled before being passed to the compiler, so if your list has many entries it could lead to very slow compilation.

An alternative would be to design your own concatenated storage format for the ragged data; say one array data with the concatenated results, and another array start_indices that indicates where each batch occurs. Then depending on what operations you wanted to do over the data, you could vmap or scan over the start_indices and operate on the particular slice in question. There has been some experimental work on this type of approach in the past; see for example the test case here: https://github.com/google/jax/blob/f1ea67117e8f102412657c513be13047d70491e9/tests/dynamic_api_test.py#L1495
But note that these Jumble APIs remain incomplete, private, and undocumented.

Finally, a third approach might be to pad your ragged arrays so they fit into a two-dimensional array, with the lengths stored in a companion array. There are tradeoffs here with respect to a concatenated approach in terms of memory and computation, but the benefit is that it would allow you to use vmap or scan directly over the leading axis.

As you can see, unfortunately there's no general one-size-fits-all approach here, but one of those three options should probably unblock you. What do you think?

[renecotyfanboy]

Thank you very much for this very detailed answer, I ended up using the start_indices trick which worked in my case. These Jumble apis were appealing, but unfortunately it seems to be too much of an experimental feature yet, but I would be glad to use it when (if) it is released!

[cottrell]

Thank you very much for this very detailed answer, I ended up using the start_indices trick which worked in my case. These Jumble apis were appealing, but unfortunately it seems to be too much of an experimental feature yet, but I would be glad to use it when (if) it is released!

Do you have an example of this? I've tried things like this

import jax
import numpy as np
import jax.numpy as jnp
from jax import lax

# Generate ragged arrays
sizes = np.random.randint(1, 4, size=16)
ragged_arrays = [np.random.randn(size) for size in sizes]

concatenated_data = jnp.concatenate(ragged_arrays)
lengths = [len(row) for row in ragged_arrays]
start_indices = jnp.cumsum(jnp.array([0] + lengths[:-1]))
lengths = jnp.array(lengths)


def scan_fun(carry, x):
    start, length = x
    slice_data = lax.dynamic_slice(concatenated_data, (start,), (length,))  # it complains about this
    s = jnp.sum(slice_data)
    return carry + s, s

coords = jnp.vstack([start_indices, lengths]).T

res = lax.scan(scan_fun, 0., coords)
# TypeError: Shapes must be 1D sequences of concrete values of integer type, got (Traced<ShapedArray(int64[])>with<DynamicJaxprTrace(level=1/0)>,).
# If using `jit`, try using `static_argnums` or applying `jit` to smaller subfunctions.

[renecotyfanboy]

I had this kind of approach :

import numpy as np
import jax.numpy as jnp
from jax import lax

# Generate ragged arrays
sizes = np.random.randint(1, 4, size=16)
ragged_arrays = jnp.asarray([np.hstack([np.random.randn(size), np.asarray([np.nan]*(3-size))]) for size in sizes])

def scan_fun(carry, x):
    array, end_idx = x
    s = jax.lax.fori_loop(0, end_idx, lambda i, s: s + array[i], 0.)
    return carry + s, s

sum_of_arrays, _ = lax.scan(scan_fun, 0., (ragged_arrays, sizes))
assert sum_of_arrays == np.nansum(ragged_arrays)

Using a fori loop instead of a dynamic slice. But it had a lot of issues, especially when nesting with grads and stuff like this

[cottrell]

I had this kind of approach :

import numpy as np
import jax.numpy as jnp
from jax import lax

# Generate ragged arrays
sizes = np.random.randint(1, 4, size=16)
ragged_arrays = jnp.asarray([np.hstack([np.random.randn(size), np.asarray([np.nan]*(3-size))]) for size in sizes])

def scan_fun(carry, x):
    array, end_idx = x
    s = jax.lax.fori_loop(0, end_idx, lambda i, s: s + array[i], 0.)
    return carry + s, s

sum_of_arrays, _ = lax.scan(scan_fun, 0., (ragged_arrays, sizes))
assert sum_of_arrays == np.nansum(ragged_arrays)

Using a fori loop instead of a dynamic slice. But it had a lot of issues, especially when nesting with grads and stuff like this

Ok so you are padding. I was hoping there might be a non-padded version as this can be costly for highly skewed sparse observations. Maybe it's a reasonably compromise as we are likely bounded by batch size on the first dimension.