code: def train_step(params, counts, batch):
      x, y = batch
  
      def loss_fn(params):
          y_pred, (updates, _) = modeldef.apply(params, counts)(x)
          counts_ = updates.extract(Count)
          loss = jnp.mean((y - y_pred) ** 2)
          return loss, counts_
  
      grad, counts = jax.grad(loss_fn, has_aux=True)(params)
      #                          |-------- sgd ---------|
      params = jax.tree_map(lambda w, g: w - 0.1 * g, params, grad)
  
      return params, counts
  
file_path: flax/flax/experimental/nnx/examples/toy_examples/01_functional_api.py
imports: import jax
  import jax.numpy
  import matplotlib.pyplot
  import numpy
  from flax.experimental import nnx
label_desc: Use image node_id 2 for calling a global function with example usage: train_step(params, counts, batch) and returns: params,  counts, loss,  counts_