Error: vmapped nnx.Module initialization with selective Variable broadcasting

[rademacher-p]

I'm trying to use nnx.vmap to create a module with different nnx.Variable attributes, only vectorizing the initialization of certain ones. The failing example below splits and vmaps over a nnx.Rngs arg, then generates two variables. I attempt to use StateAxes to only vectorize the MyVar, keeping the Param as a scalar.

Two commented lines provide modifications that succeed, but are not quite what I want.

• Mod state_axes to broadcast the Param as well
• Mod self.param to be deterministic, not needing the Rngs

import jax
import jax.numpy as jnp
from flax import nnx


class MyVar(nnx.Variable):
    pass


state_axes = nnx.StateAxes({(nnx.RngState, MyVar): 0, ...: None})
# state_axes = nnx.StateAxes({(nnx.RngState, nnx.Param, MyVar): 0, ...: None})


class MyModule(nnx.Module):
    @nnx.split_rngs(splits=2)
    @nnx.vmap(in_axes=(state_axes, 0))
    def __init__(self, rngs):
        self.param = nnx.Param(jax.random.uniform(rngs()))
        # self.param = nnx.Param(jnp.float32(0.0))
        self.var = MyVar(jax.random.uniform(rngs()))


rngs = nnx.Rngs(123)
model = MyModule(rngs)

The error is ValueError: at vmap out_axes[0][0].states[1]['param'].value, got axis spec None but output was batched on axis 0. I've tried using StateAxes for the vmap out_axes arg, to no avail. I've also tried different Rngs streams for each variable type.

Any suggestions, anyone?

@cgarciae Shameless ping. Seems like You ~= NNX 😄

[rademacher-p]

I'm guessing I don't need that RngState filter in the StateAxes, either. Might be from some other thing I tried...

[cgarciae]

Hey @rademacher-p, you need two different RNG keys, one that initializes Params as broadcasts, and one that initializes MyVars vectorized. Easiest way to do this is to use named streams in Rngs, create a params stream and a vars stream, tell split_rngs to only split the vars keys, and use the stream names to sample some keys. Here's a running sample:

import jax
import jax.numpy as jnp
from flax import nnx

class MyVar(nnx.Variable):
  pass

state_axes = nnx.StateAxes({(MyVar, 'vars'): 0, ...: None})

class MyModule(nnx.Module):
  @nnx.split_rngs(splits=2, only='vars')
  @nnx.vmap(in_axes=(state_axes, state_axes))
  def __init__(self, rngs):
    self.param = nnx.Param(jax.random.uniform(rngs.params()))
    # self.param = nnx.Param(jnp.float32(0.0))
    self.var = MyVar(jax.random.uniform(rngs.vars()))

rngs = nnx.Rngs(params=1, vars=2)
model = MyModule(rngs)

Also updated state_axes to only vectorize the vars RNG state.

[rademacher-p]

Arg, just got there myself, but you beat me :) My fix is below. Can you comment on the slightly different ways we did the StateAxes? I did separate ones for the module and the rngs arg, assuming the type filter and the stream filter are only pertinent to the module and Rngs, respectively.

import jax
from flax import nnx


class MyVar(nnx.Variable):
    pass


state_axes_self = nnx.StateAxes({MyVar: 0, ...: None})
state_axes_rngs = nnx.StateAxes({"vars": 0, ...: None})

class MyModule(nnx.Module):
    @nnx.split_rngs(splits=2, only="vars")
    @nnx.vmap(in_axes=(state_axes_self, state_axes_rngs))
    def __init__(self, rngs):
        self.param = nnx.Param(jax.random.uniform(rngs.params()))
        self.var = MyVar(jax.random.uniform(rngs.vars()))


rngs = nnx.Rngs(params=123, vars=456)
model = MyModule(rngs)

[cgarciae]

Nice! I was just lazy and reused the same StateAxes for both 😅 .

[cgarciae]

Sometime's you'll have Rngs inside Modules so you end up using a single StateAxes definition.

[rademacher-p]

Haha, that's helpful! Just made the mods on my non-toy module with an Rngs attribute, and indeed a single StateAxes suffices. Thanks for the help!