nnx.vmap use the same random key rngs inside nnx.Module across vectorization.

[maxencefaldor]

Hi,

If I run something like:

def inference_fn(model, x):
	y = model(input)
	return y

ys = nnx.vmap(inference_fn, in_axes=(None, 0))(model, xs)

the random key that is used across the vectorization is unique. It means that stochastic functions will have the same behavior across the batch.

• Is this the intended behavior for nnx.vmap when dealing with stochastic models?
• If yes, is there a recommended way to ensure independent random keys are used for each batch element when using nnx.vmap?

Thanks!

[cgarciae]

Hey, in NNX model random state is just a regular type of state, Rngs internally holds RngKey and RngCount which are subtypes of RngState(Variable). Starting from flax>=0.9.0 NNX doesn't treat random state in a special way (see JAX-style NNX Transforms), to implement RNG state handling you can either split the rng keys passed to Rngs or use the new nnx.split_rngs API (easier):

state_axes = nnx.StateAxes({RngState: 0, ...: None}) # vectorize RngState, broadcast the rest

@nnx.split_rngs(splits=<num_splits>)
@nnx.vmap(in_axes=(state_axes, 0))
def inference_fn(model, x):
	y = model(input)
	return y

split_rngs will temporarily lift/split the RngState and lower it afterwards. Note that instead of using None to broadcast the model state, you have to use StateAxes to specify that you want to vectorize the RngState on axis 0 and broadcast all other state.