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]

EDIT: Updating to use the new APIs.

Hey @JeyRunner! You can use nnx.split_rngs to automatically split the Rngs before going into nnx.vmap.

@nnx.split_rngs(splits=5)
@nnx.vmap
def make_model(rngs):
  return nnx.Linear(2, 3, rngs=rngs)

model = make_model(nnx.Rngs(0))

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.

[errhernandez]

Hi @cgarciae. In the example that you gave above, namely

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)

how would you do it if you wanted the dimensions of Linear in make_model to be input?

[cgarciae]

Hi @errhernandez, you define new arguments to make_model (e.g. din and dout) and mark them as broadcasts by setting them to None in use in_axes:

@nnx.split_rngs(splits=5)
@nnx.vmap(in_axes=(None, None, 0))
def make_model(din, dout, rngs):
  return nnx.Linear(din, dout, rngs=rngs)

[errhernandez]

Thank you @cgarciae for your quick response; your solution works, but now imagine that one wants to have variable dimensions in the stacked nnx.Linear layers; I tried what to me seemed like the obvious extension of your solution, namely

@nnx.split_rngs(splits=5)
@nnx.vmap(in_axes(0,0,0))
def make_model(din, dout, rngs):
      return nnx.Linear(din, dout, rngs=rngs)

where now din and dout are arrays of length 5. This 'obvious' solution doesn't work: when I try to call make_model I get an error message.

I am working under the assumption that the aim here is to construct a sequence of nnx.Linear layers, and that an input vector of the right dimensions would go through the first Linear, then the resulting input would go through the second, and so on, right? or am I misunderstanding something?

[cgarciae]

Hey @errhernandez, you cannot have Arrays with variable dimensions in JAX. Best you could do is have a list of models, each with different dimension, but you cannot vmap over this.

[errhernandez]

I feel like I didn't explain myself clearly: it is not that the arrays have variable dimensions; each of the elements in array din would be an input size to the corresponding nnx.Linear, and likewise each of the elements of dout would be the output size; and of course one would need to ensure that dout[i] = din[i+i], but this is the typical situation one has in an MLP, except for the missing activation functions, right?
So, are you implying that if I construct an MLP like a list of nnx.Linear and activation functions, it is not possible to vmap through it? Can you hint at what would be the way to construct an efficient MLP using nnx?
Sorry for asking so many questions...