def __call__(
      self,
      train_state: TrainState,
      **kwargs,
  ) -> tuple[TrainState, AuxOutPyTree]:
    """Signature for a training step.

    Args:
      train_state: The current state.
      **kwargs: Arguments passed to the train step, usually inputs to the model
        or labels.

    Returns:
      A tuple ``(new_train_state, aux_outputs)``.
    """
Signature for a training step.

    Args:
      train_state: The current state.
      **kwargs: Arguments passed to the train step, usually inputs to the model
        or labels.

    Returns:
      A tuple ``(new_train_state, aux_outputs)``.
    
penzai/deprecated/v1/toolshed/basic_training.py