code: def test_jit(self):
      m = nnx.Dict(a=nnx.Param(1))
  
      @nnx.jit
      def g(m: nnx.Dict):
          m.a = 2
          return 1.0
  
      out = g(m)
  
      assert m.a == 2
      assert out == 1.0
  
imports: import typing
  from functools import partial
  import jax
  import jax.numpy
  import pytest
  from flax.experimental import nnx
label_desc: Use image node_id 1 for calling the TestJIT obj's underlying member method code with example usage: obj.test_jit() and returns: int