How to Use vmap with Variable-Length Sequences in JAX?

[ajithmoola]

I'm working with JAX and trying to vectorize a function that operates on variable-length sequences derived from indices. The sequences vary in length according to a previously calculated array of segment lengths. I'm facing challenges with JAX's static shape requirements when trying to use vmap for these variable-length sequences. Here is a minimal reproducible example:

import jax.numpy as jnp
import jax
from jax import random, vmap

key = random.PRNGKey(0)

# Generate some synthetic data
n_pts = 100000
params = random.uniform(key, (n_pts, 2))
segment_lengths = random.randint(key, (n_pts,), 1, 10)
cumsum = jnp.concatenate([jnp.array([0]), segment_lengths.cumsum()])
coeff = random.normal(key, (100, 3, 3))
ind = random.randint(key, (cumsum[-1],), 0, 100)

def func(param, i, cumsum, coeff, ind):
    a = cumsum[i]
    b = cumsum[i+1]
    curr_ind = ind[jnp.arange(a, b)]
    repeated_matrix = jnp.tile(param, (b - a, 1))
    output = jnp.sum(repeated_matrix * coeff[curr_ind], axis=0)
    return output

# Attempting to apply vmap
output = jax.vmap(func, in_axes=(0, 0, None, None, None))(params, jnp.arange(n_pts), cumsum, coeff, ind)

curr_ind = ind[jnp.arange(a, b)] results in the following error:

ConcretizationTypeError: Abstract tracer value encountered where concrete value is expected: traced array with shape int32[]. It arose in the jnp.arange argument 'start'.

repeated_matrix = jnp.tile(param, (b - a, 1)) results in the following error:

TracerArrayConversionError: The numpy.ndarray conversion method array() was called on traced array with shape int32[].

How can I use vmap to handle these operations for variable-length sequences without introducing padding or restructuring my arrays significantly?

[jakevdp]

Unfortunately, it's not possible to do what you're asking with jax.vmap. We've had some experiments along these lines (see e.g. #16541 and related work) but nothing that's yet ready to use.

I suspect your best option here would be to pad all batches to the same length, and then use vmap on the padded version.

[ajithmoola]

Thank you for the response. This is a little out of scope for this discussion topic but I'm using vmap here instead of manual vectorization (I can vectorize this using scatter_add or Array.at[ind].add(arr) for all param values) but the issue is when I compute Jacobian for manually vectorized function which takes params as input, the Jacobian is huge (params and output have same shapes, If params has a shape of (N, 3), the shape of Jacobian would be (N, 3, 3, N) I guess? but vmaped function would give me a Jacobian of shape (N, 3, 3)). Is there a way to achieve this reduced Jacobian shape without using vmap?

[jakevdp]

You can slice out the part of the jacobian that you need, and wrap the whole thing in JIT. Often in such cases, the compiler is able to optimize-away the additional dimensions.

[ajithmoola]

Thanks!

[ajithmoola]

I indexed the output and wrapped everything in jax.jit but I don't think the compiler is optimizing the intermediate computations w.r.t. the indexed output. I'm getting an OOM error. @jakevdp

[jakevdp]

If you can share a minimal reproducible example of the code you're running, we may be able to help. Thanks!