code: def train_step(model: MLP, batch):
      x, y = batch
  
      def loss_fn(model: MLP):
          y_pred = model(x)
          return jnp.mean((y - y_pred) ** 2)
  
      #                                   |--default--|
      grad: nnx.State = nnx.grad(loss_fn, wrt=nnx.Param)(model)
      # sdg update
      model.update(
          jax.tree_map(
              lambda w, g: w - 0.1 * g, model.extract(nnx.Param), grad
          )
      )
  
file_path: flax/flax/experimental/nnx/examples/toy_examples/02_lifted_transforms.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(model, batch) and returns: jnp