Training state and plans

This page documents the explicit state and compiled-step contracts. Most applications can use Trainer.fit() and do not need to construct these objects directly. See the training guide for the common workflow and the Trainer reference for runtime options.

State

TrainState contains model parameters, Optax state, balancer state, independent PRNG streams, the optimizer step, and mixed-precision loss-scaling state. A complete checkpoint restores every field.

TrainState [source]

Carry all mutable arrays through compiled training computations.

Attributes:

Name	Type	Description
`model_state`	`Any`	Explicit differentiable NNX model state.
`optimizer_state`	`OptState`	Optax transformation state aligned with `model_state`.
`balancer_state`	`Any`	Arbitrary JAX-compatible loss-balancer state.
`rng_key`	`Array`	Explicit PRNG stream for stochastic model execution.
`sampling_key`	`Array`	Persistent PRNG stream for DataModule-owned training batches.
`balancer_key`	`Array`	Persistent PRNG stream for adaptive-balancer diagnostics.
`step`	`Array`	Scalar integer training iteration stored on device.
`loss_scale`	`Array`	Current loss scale used by FP16 mixed precision.
`finite_steps`	`Array`	Consecutive finite-gradient update count used to grow `loss_scale`.

initialize_train_state [source]

initialize_train_state(
    model_state: Any,
    optimizer: GradientTransformation,
    balancer_state: Any,
    key: Array,
    precision: PrecisionName | PrecisionPolicy = "32-true",
    *,
    sampling_key: Array | None = None,
    balancer_key: Array | None = None,
) -> TrainState

Initialize the complete functional training state.

Parameters:

Name	Type	Description	Default
`model_state`	`Any`	Explicit differentiable model parameter PyTree.	required
`optimizer`	`GradientTransformation`	Optax transformation used to initialize optimizer slots.	required
`balancer_state`	`Any`	Initial loss-balancer weights and diagnostics.	required
`key`	`Array`	Root training key, or the model-runtime key when both explicit persistent keys are supplied.	required
`precision`	`PrecisionName \| PrecisionPolicy`	Precision policy controlling parameter dtype and initial loss scaling.	`'32-true'`
`sampling_key`	`Array \| None`	Optional explicit DataModule sampling key.	`None`
`balancer_key`	`Array \| None`	Optional explicit adaptive-balancer diagnostic key.	`None`

Returns:

Type	Description
`TrainState`	Training state with initialized optimizer slots and device step `0`.

Raises:

Type	Description
`ValueError`	If only one explicit persistent key is supplied.

Compiled update

make_train_step() creates one JIT-compiled numerical update. It evaluates the objective, combines its losses, computes gradients, applies the Optax transformation, and returns the updated state and metrics.

make_train_step [source]

make_train_step(
    module: BasePhiModule,
    balancer: Any,
    optimizer: GradientTransformation,
    precision: PrecisionName | PrecisionPolicy = "32-true",
) -> Callable[
    [TrainState, NamedBatches],
    tuple[TrainState, dict[str, jax.Array]],
]

Create one JIT-compiled loss, gradient, and Optax update.

The returned function is created once and accepts new batch values without retracing as long as shapes, dtypes, and PyTree structures remain unchanged.

Parameters:

Name	Type	Description	Default
`module`	`BasePhiModule`	Module producing named unweighted scalar losses from explicit model state and batches.	required
`balancer`	`Any`	Balancer exposing stable `loss_names` and `combine(losses, state)`.	required
`optimizer`	`GradientTransformation`	Initialized-compatible Optax gradient transformation.	required
`precision`	`PrecisionName \| PrecisionPolicy`	Precision policy controlling batch dtype and FP16 dynamic loss scaling.	`'32-true'`

Returns:

Type	Description
`Callable[[TrainState, NamedBatches], tuple[TrainState, dict[str, Array]]]`	JIT-compiled function mapping a :class:`TrainState` and fixed-structure named batches to updated state and
`Callable[[TrainState, NamedBatches], tuple[TrainState, dict[str, Array]]]`	named scalar metrics and fixed-shape diagnostics.

Fit results and plans

FitResult returns the bound module and final functional state. TrainingPlan describes an explicit compiled update, loss balancer, and optional adaptive-balancer update. These contracts power Trainer.fit_state().

FitResult [source] `dataclass`

FitResult(
    module: BasePhiModule,
    state: TrainState,
    metrics: dict[str, float],
    stopped_early: bool,
    iterations: int,
    interrupted: bool = False,
)

Summarize the terminal state of one Trainer fit call.

Attributes:

Name	Type	Description
`module`	`BasePhiModule`	Bound module used for training and subsequent prediction.
`state`	`TrainState`	Final functional training state.
`metrics`	`dict[str, float]`	Final host scalar metrics.
`stopped_early`	`bool`	Whether a callback requested termination.
`iterations`	`int`	Number of batches processed by this fit call.
`interrupted`	`bool`	Whether an operating-system signal or :class:`KeyboardInterrupt` stopped training.

TrainingPlan [source] `dataclass`

TrainingPlan(
    train_step: TrainStep,
    batch_keys: tuple[str, ...] = (),
    balancer_update: BalancerUpdatePlan | None = None,
)

Bind a compiled optimizer step to its data and adaptive-update requirements.

Attributes:

Name	Type	Description
`train_step`	`TrainStep`	Compiled state-and-batch optimizer update that increments `TrainState.step` exactly once.
`batch_keys`	`tuple[str, ...]`	Stable named batches requested from a DataModule.
`balancer_update`	`BalancerUpdatePlan \| None`	Optional host-scheduled adaptive-balancer update.

__post_init__ [source]

__post_init__() -> None

Validate the immutable plan structure.

Raises:

Type	Description
`TypeError`	If `train_step` is not callable.
`ValueError`	If batch keys are empty or duplicated.

build_training_plan [source]

build_training_plan(
    trainer: Trainer,
    module: BasePhiModule,
    balancer: LossBalancer,
    optimizer: GradientTransformation,
) -> TrainingPlan

Resolve a configuration-independent compiled training plan.

Parameters:

Name	Type	Description	Default
`trainer`	`Trainer`	Trainer supplying precision-aware step compilation.	required
`module`	`BasePhiModule`	Bound module exposing stable loss and batch names.	required
`balancer`	`LossBalancer`	Functional loss balancer aligned with the module losses.	required
`optimizer`	`GradientTransformation`	Optax gradient transformation used by the update.	required

Returns:

Type	Description
`TrainingPlan`	Compiled optimizer step, inferred batch routing, and optional adaptive update plan.

Raises:

Type	Description
`TypeError`	If the balancer does not satisfy :class:`phijax.LossBalancer`.
`ValueError`	If loss names differ or batch names are unavailable.

Training output

TrainingOutput carries named scalar losses and optional fixed-structure diagnostics from the compiled step. Scalar diagnostics are logged by default. Array diagnostics remain available to callbacks and are never reduced implicitly.

TrainingOutput [source] `dataclass`

TrainingOutput(
    losses: Mapping[str, Array],
    diagnostics: Mapping[str, Array] = dict(),
)

Return named objective losses and optional compiled diagnostics.

Attributes:

Name	Type	Description
`losses`	`Mapping[str, Array]`	Stable scalar losses consumed by the configured loss balancer.
`diagnostics`	`Mapping[str, Array]`	Stable auxiliary arrays returned to the host without affecting gradients.

Training state and plans

State

TrainState [source]

initialize_train_state [source]

Compiled update

make_train_step [source]

Fit results and plans

FitResult [source] dataclass

TrainingPlan [source] dataclass

__post_init__ [source]

build_training_plan [source]

Training output

TrainingOutput [source] dataclass

FitResult [source] `dataclass`

TrainingPlan [source] `dataclass`

TrainingOutput [source] `dataclass`