Encountered TraceContextError after updating to 0.10.5

[zhengzl18]

Hi. My code worked fine under 0.10.4, but after updating to 0.10.5, I got an TraceContextError. My code structure was like:

class Backbone(nnx.Module):
    @nnx.jit
    def get_feature(self, obs: Obs) -> jnp.ndarray:
        ...
        return feature

class Algorithm:
    @nnx.jit
    def update(key, data, optimizer):
        def loss_fn(model):
            ...
            model.get_feature(data.obs)
            ...
        grad = nnx.grad(loss_fn)(self.model)
        optimizer.update(grad)

And the error occured in model.get_feature(data.obs):

flax.errors.TraceContextError: Cannot mutate Param from a different trace level (https://flax.readthedocs.io/en/latest/api_reference/flax.errors.html#flax.errors.TraceContextError)

I tried to not jit get_feature, but again an error occured in the dropout module:

  File "network.py", line 183, in __call__
    x = x + self.dropout(mha_out)
            ^^^^^^^^^^^^^^^^^^^^^
  File "flax/nnx/nn/stochastic.py", line 121, in __call__
    rng = rngs[self.rng_collection]()
          ^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "flax/nnx/rnglib.py", line 65, in __call__
    self._check_valid_context(
  File "flax/nnx/object.py", line 204, in _check_valid_context
    raise errors.TraceContextError(error_msg())
flax.errors.TraceContextError: Cannot call RngStream from a different trace level (https://flax.readthedocs.io/en/latest/api_reference/flax.errors.html#flax.errors.TraceContextError)

I guess it's because trace-level detection was added back in #4527? I currently downgrade back to 0.10.4, but I'm worrying about whether this indicates that something's wrong in my codes?

[jolyda]

Really odd this bug isn’t discussed anywhere at all, even by the flax maintainers.

[cgarciae]

Hey @zhengzl18 and @jolyda,

Sorry for the delay. I believe we re-enabled the mechanism that tries to prevent tracer leakage in this version. In general you should not pass mutable Flax objects as a capture. Its a bit hard to tell what is wrong without a reproducer. The only thing I notice is that Algorithm.update doesn't use self but I assume this is an error when porting it here.