Parallelizing list of nnx models on list of inputs without for loop

[uvikas]

i have two same-size lists models: List[nnx.Module] and inps: List[jax.Array]. models are separate instantiations of the same nnx.Module subclass.

i calculate outputs as: outs = [m(i) for m, i in zip(models, inps)].

how can i parallelize this without a for loop?

from what i understand, jax.vmap and jax.lax.map aren't made to do this cleanly

[IvyZX]

If your input is a python list, you have to use a for loop.

To use jaxy things like jax.vmap (which I think is what you need here), you have to squeeze them into one model and one input with an additional dimension. Example like:

linears = [nnx.Linear(8, 8, rngs=nnx.Rngs(i)) for i in range(10)]
inputs = [jax.random.normal(jax.random.key(i), (4, 8)) for i in range(10)]

linears = jax.tree.map(lambda *xs: jnp.stack(xs), *linears)
inputs = jax.tree.map(lambda *xs: jnp.stack(xs), *inputs)

@jax.vmap
def f(m, x):
  return m(x)
outs = f(linears, inputs)

[IvyZX]

If the conversion is costly, can also try to initialize with the additional dim from the beginning:

@jax.vmap
def make_linears(key):
  return nnx.Linear(8, 8, rngs=nnx.Rngs(key))
linears = make_linears(jax.random.split(jax.random.key(0), 10))