def compute_loss(
      current_params,
  ) -> tuple[jax.Array, tuple[LossStatePyTree, AuxOutPyTree]]:
    """Computes loss, aux out, and loss state for a single example."""
    # Rebuild the model using the input parameters, which we are
    # differentiating with respect to.
    model = (
        pz.select(model_without_params)
        .at_instances_of(pz.nn.Parameter)
        .apply(lambda p: dataclasses.replace(p, value=current_params[p.name]))
    )
    loss, new_loss_fn_state, aux_outputs = loss_fn(
        model=model, state=loss_fn_state, rng=step_rng, **loss_kwargs
    )
    return loss, (new_loss_fn_state, aux_outputs)
Computes loss, aux out, and loss state for a single example.
penzai/deprecated/v1/toolshed/basic_training.py