code: def test_functional_example(self):
      class Count(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 = Count(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
  
      params, counts, graphdef = model.split(nnx.Param, Count)
  
      @jax.jit
      def train_step(params, counts, x, y):
          def loss_fn(params):
              y_pred, (updates, _) = graphdef.apply(params, counts)(x)
              loss = jax.numpy.mean((y_pred - y) ** 2)
              return loss, updates.extract(Count)
  
          # compute gradient
          grads, counts = jax.grad(loss_fn, has_aux=True)(params)
          # SGD update
          params = jax.tree_map(lambda w, g: w - 0.1 * g, params, grads)
  
          return params, counts
  
      # execute the training step
      params, counts = train_step(params, counts, x, y)
      model = graphdef.merge(params, counts)
      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 4 for calling the TestIntegration obj's underlying member method code with example usage: obj.test_functional_example() and returns: params,  counts, unknown, loss,  updates