def _stateful_trainer_step(
    root_rng: PRNGKeyArray,
    state: InternalTrainerState,
    stateless_model: ModelPyTree,
    trainable_params: tuple[pz.ParameterValue, ...],
    frozen_params: tuple[pz.ParameterValue, ...],
    state_vars: tuple[pz.StateVariableValue, ...],
    loss_fn: LossFunction,
    optimizer_def: optax.GradientTransformation,
    kwargs: dict[str, Any],
) -> tuple[
    AuxOutPyTree,
    tuple[pz.ParameterValue, ...],
    tuple[pz.StateVariableValue, ...],
    InternalTrainerState,
]:
  """Implementation of the training step for StatefulTrainer."""
  step_rng = jax.random.fold_in(root_rng, state.step)

  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,
    )

  # Take gradients.
  grad_fn = jax.grad(compute_loss_and_updates, has_aux=True)
  grads, (new_state_vars, new_loss_fn_state, aux_outputs) = grad_fn(
      trainable_params
  )

  # Update parameters.
  updates, new_opt_state = optimizer_def.update(
      grads, state.opt_state, trainable_params
  )
  new_params = optax.apply_updates(trainable_params, updates)

  return (
      aux_outputs,
      new_params,
      new_state_vars,
      InternalTrainerState(
          step=state.step + 1,
          opt_state=new_opt_state,
          loss_fn_state=new_loss_fn_state,
      ),
  )
Implementation of the training step for StatefulTrainer.
penzai/toolshed/basic_training.py