def __call__(
      self,
      *,
      model: ModelPyTree,
      state: LossStatePyTree,
      rng: PRNGKeyArray,
      **kwargs,
  ) -> tuple[jax.Array, LossStatePyTree, AuxOutPyTree]:
    """Signature for a loss function.

    Args:
      model: The structure with parameters, usually a neural network model.
      state: Arbitrary state managed by the loss function. Can be None.
      rng: A JAX PRNGKey, may be ignored.
      **kwargs: Arguments passed to the train step, usually inputs to the model
        or labels.

    Returns:
      A tuple ``(loss, new_state, aux_outputs)`` for this example. ``loss``
      should be a scalar. ``new_state`` should match the structure of ``state``.
      ``aux_outputs`` can be an arbitrary PyTree.
    """
Signature for a loss function.

    Args:
      model: The structure with parameters, usually a neural network model.
      state: Arbitrary state managed by the loss function. Can be None.
      rng: A JAX PRNGKey, may be ignored.
      **kwargs: Arguments passed to the train step, usually inputs to the model
        or labels.

    Returns:
      A tuple ``(loss, new_state, aux_outputs)`` for this example. ``loss``
      should be a scalar. ``new_state`` should match the structure of ``state``.
      ``aux_outputs`` can be an arbitrary PyTree.
    
penzai/deprecated/v1/toolshed/basic_training.py