NNX + mutable state + JIT = "Cannot mutate <module name> from a different trace level"

[dfdx]

Problem you have encountered:

nnx.jit() over a module with mutable state results in an error.

Steps to reproduce:

class Count(nnx.Variable): pass

class Counter(nnx.Module):
  def __init__(self):
    self.count = Count(jnp.array(0))

  def __call__(self):
    self.count.value += 1

counter = Counter()

counter()                   # works fine
nnx.jit(counter)()       # fails

Logs, error messages, etc:

...
File ~/.local/lib/python3.10/site-packages/flax/nnx/nnx/variables.py:232, in Variable.__setattr__(self, name, value)
    231 def __setattr__(self, name: str, value: Any) -> None:
--> 232   return self._setattr(name, value)

File ~/.local/lib/python3.10/site-packages/flax/nnx/nnx/variables.py:236, in Variable._setattr(self, name, value)
    234 def _setattr(self, name: str, value: tp.Any):
    235   if not self._trace_state.is_valid():
--> 236     raise nnx.errors.TraceContextError(
    237       f'Cannot mutate {type(self).__name__} from a different trace level'
    238     )
    240   object.__setattr__(self, name, value)

TraceContextError: Cannot mutate Count from a different trace level

Any pointers to why it may be happening or what is the trace level in this case?

System information

• OS Platform and Distribution: Ubuntu 22.04
• Versions:
  • jax 0.4.29
  • jaxlib 0.4.29
  • flax 0.8.4
• Python version: 3.10.12
• GPU/TPU model and memory: NVIDIA H100, 80Gb
• CUDA version: 12.4

Disclaimer: I get the following (seemingly unrelated) warning at the start:

The NVIDIA driver's CUDA version is 12.4 which is older than the ptxas CUDA version (12.5.40). Because the driver is older than the ptxas version, XLA is disabling parallel compilation, which may slow down compilation. You should update your NVIDIA driver or use the NVIDIA-provided CUDA forward compatibility packages.

[cgarciae]

Hey! The issue is that because you are jitting counter.__call__, self is being passed as a capture and we don't allow mutating objects outside the current transform's context. To fix this simply create a aux function that accepts counter as an input and jit that:

def increment(counter: Counter):
  counter()
  
increment(counter)                   # works fine
nnx.jit(increment)(counter)      # should work

[dfdx]

Huh! I didn't think about capturing the object. Thank you!

As a mnemonic for myself:

nnx.jit(counter)()                       # fails: JIT is applied to the instance method and thus the instance itself is captured
nnx.jit(Counter.__call__)(counter)       # OK: JIT is applied to the class method, the instance is passed _after_ JIT and is NOT captured
nnx.jit(increment)(counter)              # OK: same as above, but with a helper method