nnx.Dropout in combination with jax.lax.scan returns an error

[maxencefaldor]

System information

• OS Platform and Distribution: Linux Ubuntu 22.04
• Flax, jax, jaxlib versions (obtain with pip show flax jax jaxlib:
  • flax==0.9.0
  • jax==0.4.31
  • jaxlib=0.4.31
• Python version: 3.12.4

Problem you have encountered:

nnx.Dropout in combination with jax.lax.scan returns the following error:

TraceContextError: Cannot call RngStream from a different trace level

Steps to reproduce:

from flax import nnx
from functools import partial


class MLP(nnx.Module):

	def __init__(self, *, rngs: nnx.Rngs):
		self.linear1 = nnx.Linear(10, 256, rngs=rngs)
		self.linear2 = nnx.Linear(256, 10, rngs=rngs)
		self.dropout = nnx.Dropout(rate=0.5, rngs=rngs)

	def __call__(self, x):
		x = nnx.relu(self.dropout(self.linear1(x)))
		x = self.linear2(x)
		return x

model = MLP(rngs=nnx.Rngs(0))

model(jnp.ones((1, 10))) works fine.

However,

def scan(carry, x):
    model(x)
    return carry, None

jax.lax.scan(scan, None, jnp.ones((10, 10)))

returns an error:

---------------------------------------------------------------------------
TraceContextError                         Traceback (most recent call last)
Cell In[91], [line 5](vscode-notebook-cell:?execution_count=91&line=5)
      [2](vscode-notebook-cell:?execution_count=91&line=2)     model(x)
      [3](vscode-notebook-cell:?execution_count=91&line=3)     return carry, None
----> [5](vscode-notebook-cell:?execution_count=91&line=5) jax.lax.scan(scan, None, jnp.ones((10, 10)))

    [... skipping hidden 9 frame]

Cell In[91], [line 2](vscode-notebook-cell:?execution_count=91&line=2)
      [1](vscode-notebook-cell:?execution_count=91&line=1) def scan(carry, x):
----> [2](vscode-notebook-cell:?execution_count=91&line=2)     model(x)
      [3](vscode-notebook-cell:?execution_count=91&line=3)     return carry, None

Cell In[88], [line 13](vscode-notebook-cell:?execution_count=88&line=13)
     [12](vscode-notebook-cell:?execution_count=88&line=12) def __call__(self, x):
---> [13](vscode-notebook-cell:?execution_count=88&line=13) 	x = nnx.relu(self.dropout(self.linear1(x)))
     [14](vscode-notebook-cell:?execution_count=88&line=14) 	x = self.linear2(x)
     [15](vscode-notebook-cell:?execution_count=88&line=15) 	return x

File /src/conda/lib/python3.12/site-packages/flax/nnx/nnx/nn/stochastic.py:135, in Dropout.__call__(self, inputs, deterministic, rngs)
    [127](https://vscode-remote+ssh-002dremote-002bcax-007edogwood.vscode-resource.vscode-cdn.net/src/conda/lib/python3.12/site-packages/flax/nnx/nnx/nn/stochastic.py:127) rngs = first_from(
    [128](https://vscode-remote+ssh-002dremote-002bcax-007edogwood.vscode-resource.vscode-cdn.net/src/conda/lib/python3.12/site-packages/flax/nnx/nnx/nn/stochastic.py:128)   rngs,
    [129](https://vscode-remote+ssh-002dremote-002bcax-007edogwood.vscode-resource.vscode-cdn.net/src/conda/lib/python3.12/site-packages/flax/nnx/nnx/nn/stochastic.py:129)   self.rngs,
    [130](https://vscode-remote+ssh-002dremote-002bcax-007edogwood.vscode-resource.vscode-cdn.net/src/conda/lib/python3.12/site-packages/flax/nnx/nnx/nn/stochastic.py:130)   error_msg="""`deterministic` is False, but no `rngs` argument was provided to Dropout
    [131](https://vscode-remote+ssh-002dremote-002bcax-007edogwood.vscode-resource.vscode-cdn.net/src/conda/lib/python3.12/site-packages/flax/nnx/nnx/nn/stochastic.py:131)       as either a __call__ argument or class attribute.""",
    [132](https://vscode-remote+ssh-002dremote-002bcax-007edogwood.vscode-resource.vscode-cdn.net/src/conda/lib/python3.12/site-packages/flax/nnx/nnx/nn/stochastic.py:132) )
    [134](https://vscode-remote+ssh-002dremote-002bcax-007edogwood.vscode-resource.vscode-cdn.net/src/conda/lib/python3.12/site-packages/flax/nnx/nnx/nn/stochastic.py:134) keep_prob = 1.0 - self.rate
--> [135](https://vscode-remote+ssh-002dremote-002bcax-007edogwood.vscode-resource.vscode-cdn.net/src/conda/lib/python3.12/site-packages/flax/nnx/nnx/nn/stochastic.py:135) rng = rngs[self.rng_collection]()
    [136](https://vscode-remote+ssh-002dremote-002bcax-007edogwood.vscode-resource.vscode-cdn.net/src/conda/lib/python3.12/site-packages/flax/nnx/nnx/nn/stochastic.py:136) broadcast_shape = list(inputs.shape)
    [137](https://vscode-remote+ssh-002dremote-002bcax-007edogwood.vscode-resource.vscode-cdn.net/src/conda/lib/python3.12/site-packages/flax/nnx/nnx/nn/stochastic.py:137) for dim in self.broadcast_dims:

File /src/conda/lib/python3.12/site-packages/flax/nnx/nnx/rnglib.py:81, in RngStream.__call__(self)
     [80](https://vscode-remote+ssh-002dremote-002bcax-007edogwood.vscode-resource.vscode-cdn.net/src/conda/lib/python3.12/site-packages/flax/nnx/nnx/rnglib.py:80) def __call__(self) -> jax.Array:
---> [81](https://vscode-remote+ssh-002dremote-002bcax-007edogwood.vscode-resource.vscode-cdn.net/src/conda/lib/python3.12/site-packages/flax/nnx/nnx/rnglib.py:81)   self.check_valid_context(
     [82](https://vscode-remote+ssh-002dremote-002bcax-007edogwood.vscode-resource.vscode-cdn.net/src/conda/lib/python3.12/site-packages/flax/nnx/nnx/rnglib.py:82)     lambda: 'Cannot call RngStream from a different trace level'
     [83](https://vscode-remote+ssh-002dremote-002bcax-007edogwood.vscode-resource.vscode-cdn.net/src/conda/lib/python3.12/site-packages/flax/nnx/nnx/rnglib.py:83)   )
     [84](https://vscode-remote+ssh-002dremote-002bcax-007edogwood.vscode-resource.vscode-cdn.net/src/conda/lib/python3.12/site-packages/flax/nnx/nnx/rnglib.py:84)   key = jax.random.fold_in(self.key.value, self.count.value)
     [85](https://vscode-remote+ssh-002dremote-002bcax-007edogwood.vscode-resource.vscode-cdn.net/src/conda/lib/python3.12/site-packages/flax/nnx/nnx/rnglib.py:85)   self.count.value += 1

File /src/conda/lib/python3.12/site-packages/flax/nnx/nnx/object.py:131, in Object.check_valid_context(self, error_msg)
    [129](https://vscode-remote+ssh-002dremote-002bcax-007edogwood.vscode-resource.vscode-cdn.net/src/conda/lib/python3.12/site-packages/flax/nnx/nnx/object.py:129) def check_valid_context(self, error_msg: tp.Callable[[], str]) -> None:
    [130](https://vscode-remote+ssh-002dremote-002bcax-007edogwood.vscode-resource.vscode-cdn.net/src/conda/lib/python3.12/site-packages/flax/nnx/nnx/object.py:130)   if not self._object__state.trace_state.is_valid():
--> [131](https://vscode-remote+ssh-002dremote-002bcax-007edogwood.vscode-resource.vscode-cdn.net/src/conda/lib/python3.12/site-packages/flax/nnx/nnx/object.py:131)     raise errors.TraceContextError(error_msg())

TraceContextError: Cannot call RngStream from a different trace level

[cgarciae]

This makes sense, nnx.Dropout updates is RNG state but you cannot update of NNX objects passed as a capture to JAX transforms. To fix this use nnx.scan and pass model as a carry:

class MLP(nnx.Module):
  def __init__(self, *, rngs: nnx.Rngs):
    self.linear1 = nnx.Linear(10, 256, rngs=rngs)
    self.linear2 = nnx.Linear(256, 10, rngs=rngs)
    self.dropout = nnx.Dropout(rate=0.5, rngs=rngs)

  def __call__(self, x):
    x = nnx.relu(self.dropout(self.linear1(x)))
    x = self.linear2(x)
    return x

model = MLP(rngs=nnx.Rngs(0))

@nnx.scan  # nnx.scan has similar api to vmap
def scan_fn(model, x):
  y = model(x)
  return model, y

model, y = scan_fn(model, jnp.ones((10, 10)))

You will get new RNG state each step.