nnx using vmap to create multiple models

[JeyRunner]

Hi, I want to create multiple instances of a nnx.module (each initialized with a different key).
With linen, I can do:

def make(rng):
  m = my_module.init(rng, dummy_input)
  return ...

rngs = jax.random.split(jax.random.PRNGKey(0), num=5)
models = jax.vmap(make)(rngs)

How can I achieve the same with nnx?
I tried:

def make_model(rngs):
  return nnx.Sequential(
    nnx.Linear(..., rngs=rngs),
    ...
  )

init_keys = jax.random.split(jax.random.PRNGKey(0), num=5)
rngs = nnx.Rngs(init_keys)

model = jax.vmap(task.make_model)(rngs)

But I get the error TypeError: object of type 'RngStream' has no len().

[cgarciae]

Hey @JeyRunner! nnx.vmap has some support splitting RNG state automatically, try:

def make_model(rngs):
  return nnx.Linear(2, 3, rngs=rngs)

rngs = nnx.Rngs(0)
model = nnx.vmap(make_model, split_rngs=True, axis_size=5)(rngs)

print(model)

Output:

Linear(
  bias=Param(
    value=Array(shape=(5, 3), dtype=float32)
  ),
  bias_init=<function zeros at 0x11ee95f30>,
  dot_general=<function dot_general at 0x11e933910>,
  dtype=None,
  in_features=2,
  kernel=Param(
    value=Array(shape=(5, 2, 3), dtype=float32)
  ),
  kernel_init=<function variance_scaling.<locals>.init at 0x11fa8fe20>,
  out_features=3,
  param_dtype=<class 'jax.numpy.float32'>,
  precision=None,
  use_bias=True
)

[JeyRunner]

Thanks @cgarciae, thats exactly what I wanted 😃 !
I didn't know that nxx.vmap can do that.

[JeyRunner]

In the meantime I did this:

def make_model(rng):
  rngs = nnx.Rngs(rng)
  return nnx.Linear(2, 3, rngs=rngs)

init_keys = jax.random.split(jax.random.PRNGKey(0), num=5)
model = jax.vmap(task.make_model)(init_keys)

This should do kind of the same thing right?

[cgarciae]

You must either use NNX transforms (nnx.vmap, nnx.jit, nnx.grad, etc) to transform functions that accept Modules or use the Functional API and handle state manually.