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...

[BeeGass]

@cgarciae Im trying to understand how to use vmap such that I have multiple models for more than just one dimensions axis. Lets say that we have data that we have restructured to have the shape:
B, H, T, DH = x.shape

where:
B: batch
H: heads
T: sequence length
HD: head embed dimension

where we want to parallelize over both the batch dimension (B) and also the head dimension (H). Is this a question of nesting vmaps such that:

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

@nnx.split_rngs(splits=d_head)
@nnx.vmap(in_axes=0)
def make_d_head_vmap_model(rngs):
    return make_batch_vmap_model(rngs)

perhaps this is completely incorrect and that I actually should be using nnx.pmap

Thanks in advance for the help

[cgarciae]

Hey @BeeGass sorry for the late replay.

It works if you specify in_axes=0 on the last vmap:

num_devices = 8
per_device_batch = 4
num_heads = 2
din = 16
dout = 32

@nnx.split_rngs(splits=num_devices)
@nnx.pmap(in_axes=0, out_axes=0)  # device dim
@nnx.split_rngs(splits=per_device_batch)
@nnx.vmap(in_axes=0, out_axes=0)  # batch dim
@nnx.split_rngs(splits=num_heads)
@nnx.vmap(in_axes=0, out_axes=1)  # head dim
def make_model(rngs):
  return nnx.Linear(din, dout, rngs=rngs)

model = make_model(nnx.Rngs(0))
print(model)

Linear(
  bias=Param(
    value=Array(shape=(8, 4, 32, 2), dtype=float32)
  ),
  bias_init=<function zeros at 0x137c5bac0>,
  dot_general=<function dot_general at 0x127c03a30>,
  dtype=None,
  in_features=16,
  kernel=Param(
    value=Array(shape=(8, 4, 16, 2, 32), dtype=float32)
  ),
  kernel_init=<function variance_scaling.<locals>.init at 0x142594c10>,
  out_features=32,
  param_dtype=<class 'jax.numpy.float32'>,
  precision=None,
  use_bias=True
)

[cgarciae]

Reason is that the incoming keys on the Rngs objects only have one dimension so you cannot vectorize on axis 1.

[BeeGass]

I don't quite understand this. Is this because the axes here are taking in the zeroth rng? If this is the case how does one ensure that the model is truly vectorizing over the dimensions described? Is it just then a question of reshaping the input correctly assuming that pmap/vmap takes in the zeroth input, for each?

Also not to worry at all @cgarciae . I appreciate the help.

[cgarciae]

No worries, happy to help! The reason is that split_rngs takes a key that has shape () and turns it into a key with shape (splits,), so it only has available axis 0 for vmap.

[BeeGass]

Reason is that the incoming keys on the Rngs objects only have one dimension so you cannot vectorize on axis 1.

Oh I didn't see this. That makes sense. Thank you.