code: def test_stateful_example(self):
      class State(nnx.Variable[A]):
          pass
  
      class Linear(nnx.Module):
          def __init__(self, din: int, dout: int, *, rngs: nnx.Rngs):
              key = rngs.params()
              self.w = nnx.Param(jax.random.uniform(key, (din, dout)))
              self.b = nnx.Param(jnp.zeros((dout,)))
              self.count = State(0)
  
          def __call__(self, x):
              self.count += 1
              return x @ self.w + self.b[None]
  
      model = Linear(din=12, dout=2, rngs=nnx.Rngs(0))
      # forward pass
      x = jnp.ones((8, 12))
      y = model(x)
      assert model.count == 1
  
      @nnx.jit
      def train_step(model, x, y):
          def loss_fn(model):
              y_pred = model(x)
              return jax.numpy.mean((y_pred - y) ** 2)
  
          # compute gradient
          grads: nnx.State = nnx.grad(loss_fn, wrt=nnx.Param)(model)
          # SGD update
          model.update(
              jax.tree_map(
                  lambda w, g: w - 0.1 * g,
                  model.extract(nnx.Param),
                  grads,
              )
          )
  
      # execute the training step
      train_step(model, x, y)
      assert model.count == 2
  
file_path: flax/flax/experimental/nnx/tests/test_integration.py
imports: import typing
  import jax
  import jax.numpy
  import numpy
  from flax.experimental import nnx
label_desc: Use image node_id 3 for calling the TestIntegration obj's underlying member method code with example usage: obj.test_stateful_example() and returns: unknown, jax