code: def test_trace_level(self):
      m = nnx.Dict(a=nnx.Param(1))
  
      @jax.jit
      def f():
          with pytest.raises(
              nnx.TraceContextError,
              match="Cannot mutate Module from different trace level",
          ):
              m.a = 2
  
      f()
  
imports: import dataclasses
  from copy import deepcopy
  from typing import Any, TypeVar
  import jax
  import jax.numpy
  import pytest
  from flax.experimental import nnx
label_desc: Use image node_id 2 for calling the TestModule obj's underlying member method code with example usage: obj.test_trace_level() without return types