code: def test_jit_combinator(self):
      n = 0
  
      class Foo(nnx.Module):
          def __init__(self, din: int, dout: int, *, rngs: nnx.Rngs):
              key = rngs.params()
              self.w = nnx.Param(
                  jax.random.normal(key, shape=(din, dout))
              )
              self.din = din
              self.dout = dout
  
          @nnx.jit
          def __call__(self, x: jax.Array) -> jax.Array:
              nonlocal n
              n += 1
              return jnp.dot(x, self.w)
  
      m = nnx.JIT(Foo)(2, 3, rngs=nnx.Rngs(0))
  
      y = m(jnp.ones((1, 2)))
      assert y.shape == (1, 3)
      assert n == 1
      y = m(jnp.ones((1, 2)))
      assert n == 1
  
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 4 for calling the TestJIT obj's underlying member method code with example usage: obj.test_jit_combinator() and returns: jnp