code: def test_step(state: nnx.TrainState[MLP], batch):
      x, y = batch
      y_pred, _ = state.apply("params", "counts")(x)
      loss = jnp.mean((y - y_pred) ** 2)
      return {"loss": loss}
  
file_path: flax/flax/experimental/nnx/examples/toy_examples/03_train_state.py
imports: import jax
  import jax.numpy
  import matplotlib.pyplot
  import numpy
  import optax
  from flax.experimental import nnx
label_desc: Use image node_id 3 for calling a global function with example usage: test_step(state, batch) and returns: dict