def compute_loss_and_updates(
      trainable_params: tuple[pz.ParameterValue, ...],
  ) -> tuple[
      jax.Array,
      tuple[tuple[pz.ParameterValue, ...], LossStatePyTree, AuxOutPyTree],
  ]:
    """Computes loss, aux out, and state updates for a single example."""
    # Rebuild the model using the input parameters, which we are
    # differentiating with respect to, and the current states. Parameters are
    # kept frozen, since they will be updated by the optimizer, not the model.
    state_vars_mut = tuple(v.unfreeze_as_copy() for v in state_vars)
    model = pz.bind_variables(
        stateless_model,
        state_vars_mut + trainable_params + frozen_params,
    )
    loss, new_loss_fn_state, aux_outputs = loss_fn(
        model=model, state=state.loss_fn_state, rng=step_rng, **kwargs
    )
    return loss, (
        pz.freeze_variables(state_vars_mut),
        new_loss_fn_state,
        aux_outputs,
    )
Computes loss, aux out, and state updates for a single example.
penzai/toolshed/basic_training.py