code: def test_step(model: MLP, batch):
      x, y = batch
      y_pred = model(x)
      loss = jnp.mean((y - y_pred) ** 2)
      return {"loss": loss}
  
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 3 for calling a global function with example usage: test_step(model, batch) and returns: dict