def compute_training_outputs_and_updates(
    current_params: dict[str, Any],
    model_without_params: ModelPyTree,
    opt_state: OptimizerStatePyTree,
    loss_fn_state: LossStatePyTree,
    root_rng: PRNGKeyArray,
    step: int | jax.Array,
    loss_kwargs: dict[str, Any],
    loss_fn: LossFunction,
    optimizer_def: optax.GradientTransformation,
) -> tuple[ModelPyTree, OptimizerStatePyTree, LossStatePyTree, AuxOutPyTree]:
  """Runs a loss function and computes all of its outputs.

  This function runs the model and loss function and updates the corresponding
  parameters in the optimizer. It splits each component of the input into a
  separate argument to make it easy to JIT-compile, and to allow donating the
  parts that will be updated.

  Args:
    current_params: A dictionary of model parameters. These are the parts that
      WILL be updated by the optimizer.
    model_without_params: A model PyTree that includes ``pz.NotInThisPartition``
      in place of each of the learnable parameter values. These are the parts of
      the model that will NOT be updated by the optimizer.
    opt_state: State for the optimizer.
    loss_fn_state: State for the loss function.
    root_rng: Root random key for the training process.
    step: Current step of training, used to adjust root RNG.
    loss_kwargs: Keyword arguments for the loss function.
    loss_fn: The loss function.
    optimizer_def: The optimizer.

  Returns:
    Tuple of ``(new_params, new_opt_state, new_loss_fn_state, aux_outputs)``
  """
  step_rng = jax.random.fold_in(root_rng, step)

  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)

  # Take gradients.
  grad_fn = jax.grad(compute_loss, has_aux=True)
  grads, (new_loss_fn_state, aux_outputs) = grad_fn(current_params)

  # Get updated parameters.
  updates, new_opt_state = optimizer_def.update(
      grads, opt_state, current_params
  )
  new_params = optax.apply_updates(current_params, updates)

  return new_params, new_opt_state, new_loss_fn_state, aux_outputs
Runs a loss function and computes all of its outputs.

  This function runs the model and loss function and updates the corresponding
  parameters in the optimizer. It splits each component of the input into a
  separate argument to make it easy to JIT-compile, and to allow donating the
  parts that will be updated.

  Args:
    current_params: A dictionary of model parameters. These are the parts that
      WILL be updated by the optimizer.
    model_without_params: A model PyTree that includes ``pz.NotInThisPartition``
      in place of each of the learnable parameter values. These are the parts of
      the model that will NOT be updated by the optimizer.
    opt_state: State for the optimizer.
    loss_fn_state: State for the loss function.
    root_rng: Root random key for the training process.
    step: Current step of training, used to adjust root RNG.
    loss_kwargs: Keyword arguments for the loss function.
    loss_fn: The loss function.
    optimizer_def: The optimizer.

  Returns:
    Tuple of ``(new_params, new_opt_state, new_loss_fn_state, aux_outputs)``
  
penzai/deprecated/v1/toolshed/basic_training.py