Static argument in vmapped function?

[leonard-gleyzer]

Hello! First, thank you for all the work put into JAX. I've greatly enjoyed using it. I have a question. Is there any way to do something like the following?

import jax.numpy as jnp
from jax import jit, vmap
from functools import partial

arrays = [jnp.ones((3,)), jnp.ones((5,))]

@partial(jit, static_argnums=0)
def f(i):
    return jnp.sum(arrays[i])

indices = jnp.arange(2)

vmap(f)(indices)

This throws an error because "indices" becomes a BatchTracer, which is unhashable. I would like to be able to vmap/parallelize a function (that always returns a scalar) over a collection of 1d arrays that have different numbers of elements. I know one option is to use masking, but for my specific application, where speed/performance is important, the masking approach does not scale well.

[jakevdp]

"static" arguments in jax.jit are ones that are not traced. To make an argument untraced in vmap you can set its corresponding in_axes to None.

In your case, though, it sounds like you're asking how you can map over an argument like indices without tracing over it – that is impossible with vmap by design: the mechanism by which vmap maps over arguments is by tracing them, so a mapped argument cannot be static in the sense you're asking about.

The best way to handle the particular example in your question would probably be a list comprehension:

result = [f(i) for i in indices]

[YouJiacheng]

JAX dispatch computation asynchronously on GPU, so you can simply loop over indices.