basic_training

Basic training logic for training simple models.

This module provides a barebones implementation of training logic that supports training Penzai models. This can be used to train simple models that do not require more sophisticated training loops. It also serves as a starting point for more complex training scripts.

Classes

InternalTrainerState	Internal state for the trainer.
LossFunction	Signature for loss functions expected by the common training step.
StatefulTrainer	A trainer object that updates its state in place.