Trainer

Trainer runs the host side of fitting and prediction. It prepares data sources, places arrays, calls one compiled update per step, dispatches hooks, logs metrics, saves checkpoints, handles interruption, and releases task resources.

Basic use

Construct a Trainer, fit a PhiModule blueprint, and pass the returned FitResult to predict():

from phijax import Trainer

trainer = Trainer(max_steps=10_000, accelerator="auto")
result = trainer.fit(
    module,
    datamodule=data_module,
    optimizer=optimizer,
    seed=0,
)
predictions = trainer.predict(result, datamodule=data_module)

The Trainer initializes the model and functional state during fit(). Applications do not need to call close() after a normal fit or prediction task. See Train and predict for checkpoint restoration, adaptive loss balancing, interruption, and the explicit state API.

What the Trainer manages

The Trainer owns orchestration around the compiled numerical update:

DataModule source -> cast and place batch -> compiled update -> hooks -> logging and checkpointing
                                              |
                                              +-> explicit TrainState

It manages:

• DataModule setup, named batch construction, prediction batches, and teardown;
• model initialization and independent runtime, sampling, and balancer PRNG streams;
• precision conversion and device or data-parallel placement;
• callback and module lifecycle ordering;
• metric routing, progress displays, model summaries, and experiment loggers;
• checkpoint restoration and persistence; and
• graceful interruption and resource cleanup.

Losses, derivatives, gradients, balancing, and Optax updates remain inside the JIT-compiled step. Python hooks and filesystem operations remain outside it.

Configuration

Execution

Argument	Supported values	Default
max_steps	Positive integer	Required
callbacks	Ordered iterable of callback instances	()
compilation_cache	Mapping with enabled and, when enabled, directory	None

max_steps counts optimizer updates within one fit call. A custom step passed through fit_state() must increment TrainState.step exactly once per completed update.

Deterministic GPU execution

The Trainer does not configure process-wide GPU determinism. Launch Python with the deterministic XLA options set before importing JAX:

XLA_FLAGS="--xla_gpu_exclude_nondeterministic_ops --xla_gpu_autotune_level=0" python train.py

See Randomness and reproducibility for seed requirements, performance and compilation trade-offs, and the official OpenXLA and JAX references.

Devices and precision

Argument	Supported values	Default
accelerator	"auto", "cpu", "gpu", "tpu"	"auto"
devices	Positive count, device-index sequence, or "auto"	1
strategy	A configured Strategy, which replaces accelerator and devices	None
precision	"64-true", "32-true", "16-true", "bf16-true", "16-mixed", "bf16-mixed"	"32-true"
matmul_precision	None, "default", "high", "highest"	None
derivative_dtype	Floating dtype override	None
initial_loss_scale	Positive FP16 dynamic-loss scale	32768.0
loss_scale_growth_interval	Positive interval between finite FP16 scale increases	2000

An explicit Strategy controls placement and distributed behavior. See Precision for dtype behavior and supported device combinations.

Logging and display

Argument	Supported values	Default
logger	True, False, None, one logger, or an iterable of loggers	True
default_root_dir	Parent directory for default versioned logs	"."
log_every_n_steps	Positive scalar-logging interval	10
enable_progress_bar	Boolean	True
enable_model_summary	Boolean	True

logger=True selects TensorBoard when it is installed and CSV otherwise. The Trainer adds a TQDM progress bar and a plain model summary unless callbacks provide replacements. See Loggers and Callbacks for backend and display options.

Common and advanced workflows

Use fit() and predict() for normal application code. The Trainer initializes and binds the model, constructs the training plan, creates state, and obtains batches from the DataModule.

Use fit_state() and predict_state() when an integration supplies a bound BasePhiModule, explicit TrainState, custom TrainingPlan, or custom batch source. Supporting methods such as initialize_state(), compile_train_step(), prepare_batch_source(), and prepare_batch() expose the same runtime policies without creating another training loop.

Checkpoint-enabled applications can use ckpt_path="last" with fit(), call latest_checkpoint() explicitly, or load model parameters into an existing state with load_weights().

Metrics and status

The Trainer exposes the latest metric views without attaching mutable Trainer state to a module:

Property	Contents
callback_metrics	Complete merged metric mapping available to callbacks
logged_metrics	Scalar metrics selected for experiment loggers
progress_bar_metrics	Scalar metrics selected for progress displays
interrupted	Whether the latest fit stopped after an interrupt
received_sigterm	Whether the latest fit received SIGTERM
logger	Configured logger collection

Metric mappings contain device values until a logger, display, checkpoint monitor, or caller requires host conversion. See Logging and monitoring for module and callback metric routing.

Cleanup and interruption

Each fit() and predict() call releases DataModule, callback, logger, and checkpoint resources after success, interruption, or failure. A Trainer can be reused without an explicit close() call. The idempotent close() method and context-manager support remain available for integrations that open resources outside a Trainer task.

The first Ctrl+C returns the last completed state with FitResult.interrupted=True. SIGTERM runs exception hooks and cleanup before exiting with the standard signal status. Environment information is printed when the Trainer is created and only on global rank zero.

API reference

Trainer [source]

Trainer(
    max_steps: int,
    *,
    accelerator: Accelerator = "auto",
    devices: DeviceSelection = 1,
    precision: PrecisionName | PrecisionPolicy = "32-true",
    matmul_precision: MatmulPrecision | None = None,
    derivative_dtype: Any | None = None,
    initial_loss_scale: float = 32768.0,
    loss_scale_growth_interval: int = 2000,
    log_every_n_steps: int = 10,
    enable_progress_bar: bool = True,
    enable_model_summary: bool = True,
    callbacks: Iterable[Callback] = (),
    logger: bool
    | ExperimentLogger
    | Iterable[ExperimentLogger]
    | None = True,
    default_root_dir: str | Path = ".",
    loggers: Iterable[ExperimentLogger] | None = None,
    strategy: Strategy | None = None,
    compilation_cache: Mapping[str, Any] | None = None,
)

Coordinate compiled JAX updates, callbacks, logging, and checkpointing on the host.

The trainer never owns a mutable model object and never places Python hooks inside :func:jax.jit. Numerical behavior remains in the configured training plan or supplied compiled step; this class orchestrates its lifecycle.

Attributes:

Name	Type	Description
prediction_writer		Configured prediction artifact callback, or None.

Initialize host-side training orchestration.

Parameters:

Name	Type	Description	Default
max_steps	int	Maximum number of batches processed by one fit call.	required
accelerator	Accelerator	Device backend selected when strategy is absent. Choose auto, cpu, gpu, or tpu.	'auto'
devices	DeviceSelection	Positive device count, explicit device indices, or auto for all matching visible devices.	1
precision	PrecisionName | PrecisionPolicy	Precision name or resolved policy. Supported names are 64-true, 32-true, 16-true, bf16-true, 16-mixed, and bf16-mixed; common dtype aliases are also accepted.	'32-true'
matmul_precision	MatmulPrecision | None	Optional default, high, or highest JAX matrix-multiplication precision override.	None
derivative_dtype	Any | None	Optional floating batch and coordinate-derivative dtype override.	None
initial_loss_scale	float	Initial dynamic scale for 16-mixed training.	32768.0
loss_scale_growth_interval	int	Finite optimizer updates between FP16 scale increases.	2000
log_every_n_steps	int	Positive interval between scalar logger updates.	10
enable_progress_bar	bool	Whether to enable a TQDM or explicitly configured progress callback.	True
enable_model_summary	bool	Whether to add a plain model summary when none is supplied explicitly.	True
callbacks	Iterable[Callback]	Ordered host-side lifecycle callbacks.	()
logger	bool | ExperimentLogger | Iterable[ExperimentLogger] | None	Default logger flag, one backend, several backends, or None to disable logging.	True
default_root_dir	str | Path	Parent directory for versioned default local logs.	'.'
loggers	Iterable[ExperimentLogger] | None	Deprecated plural alias for an explicit iterable of logging backends.	None
strategy	Strategy | None	Explicit placement strategy overriding accelerator and devices.	None
compilation_cache	Mapping[str, Any] | None	Optional mapping with enabled and directory values applied before device discovery.	None

Raises:

Type	Description
TypeError	If a display flag or logger configuration has an invalid type.
ValueError	If a step, interval, callback, logger, or matrix-multiplication precision option is invalid.

logger property

logger: LoggerCollection

Return the configured experiment logger collection.

Returns:

Type	Description
LoggerCollection	Active logger collection.

interrupted property writable

interrupted: bool

Return whether the latest fit was interrupted.

Returns:

Type	Description
bool	Current signal connector interruption state.

received_sigterm property writable

received_sigterm: bool

Return whether the latest fit received SIGTERM.

Returns:

Type	Description
bool	Current signal connector termination flag.

logged_metrics property

logged_metrics: Mapping[str, Any]

Return metrics selected for experiment loggers at the latest completed step.

Returns:

Type	Description
Mapping[str, Any]	Immutable-by-convention logger metric mapping.

progress_bar_metrics property

progress_bar_metrics: Mapping[str, Any]

Return metrics selected for progress display at the latest completed step.

Returns:

Type	Description
Mapping[str, Any]	Immutable-by-convention progress metric mapping.

callback_metrics property

callback_metrics: Mapping[str, Any]

Return the complete latest metric mapping used by callbacks and fit results.

Returns:

Type	Description
Mapping[str, Any]	Complete merged metric mapping.

fit [source]

fit(
    module: PhiModule,
    *,
    datamodule: PhiDataModule,
    optimizer: GradientTransformation,
    seed: int | Array,
    balancer: LossBalancer | None = None,
    hyperparameters: Mapping[str, Any] | None = None,
    ckpt_path: str | Path | None = None,
    ckpt_step: int | None = None,
    weights_only: bool = False,
) -> FitResult

Initialize and fit a module blueprint through the concise application API.

Parameters:

Name	Type	Description	Default
module	PhiModule	Unbound :class:PhiModule containing a model factory and objective.	required
datamodule	PhiDataModule	Application DataModule supplying normalization statistics and training batches.	required
optimizer	GradientTransformation	Optax transformation used to initialize state and compile updates.	required
seed	int | Array	Integer seed or valid JAX key. PhiJAX does not modify Python or NumPy global random state.	required
balancer	LossBalancer | None	Functional loss balancer, or None for equal static weights.	None
hyperparameters	Mapping[str, Any] | None	Optional resolved configuration recorded before the first update.	None
ckpt_path	str | Path | None	Optional Orbax checkpoint root, or "last" for the configured callback's latest checkpoint.	None
ckpt_step	int | None	Exact checkpoint step, or None for the latest committed step.	None
weights_only	bool	Whether to restore only model weights into the fresh state.	False

Returns:

Type	Description
FitResult	Bound module, final state, metrics, and termination status.

Raises:

Type	Description
TypeError	If seed is not an integer or JAX key.
ValueError	If model, balancing, checkpoint, or DataModule options are inconsistent.

predict [source]

predict(
    result: FitResult,
    batches: Iterable[Mapping[str, Any]] | None = None,
    *,
    datamodule: PhiDataModule | None = None,
    ckpt_path: str | Path | None = None,
    ckpt_step: int | None = None,
    metadata: Mapping[str, Any] | None = None,
    return_predictions: bool = True,
) -> np.ndarray | None

Predict from the bound module and functional state returned by :meth:fit.

Parameters:

Name	Type	Description	Default
result	FitResult	Completed fit result containing the bound module and state.	required
batches	Iterable[Mapping[str, Any]] | None	Optional finite prediction batch iterable. When omitted, use datamodule.	None
datamodule	PhiDataModule | None	Optional DataModule supplying prediction batches.	None
ckpt_path	str | Path | None	Optional Orbax checkpoint root, or "last" for the configured callback's latest checkpoint.	None
ckpt_step	int | None	Exact checkpoint step, or None for the latest committed step.	None
metadata	Mapping[str, Any] | None	Optional immutable metadata exposed to prediction callbacks.	None
return_predictions	bool	Whether to concatenate and return host predictions.	True

Returns:

Type	Description
ndarray | None	Concatenated host predictions, or None when collection is disabled or prediction data is unavailable.

fit_state [source]

fit_state(
    module: BasePhiModule,
    training: TrainingPlan | TrainStep,
    state: TrainState,
    batches: BatchSource | None = None,
    *,
    hyperparameters: Mapping[str, Any] | None = None,
    ckpt_path: str | Path | None = None,
    ckpt_step: int | None = None,
    weights_only: bool = False,
    datamodule: PhiDataModule | None = None,
) -> FitResult

Fit an explicitly initialized module, state, and resolved numerical plan.

Parameters:

Name	Type	Description	Default
module	BasePhiModule	Bound custom module owning model, objective, and lifecycle behavior.	required
training	TrainingPlan | TrainStep	Resolved training plan, or a compiled step when batches is supplied explicitly.	required
state	TrainState	Initial functional training state and restore template.	required
batches	BatchSource | None	Optional explicit step-indexed callable or finite iterable source.	None
hyperparameters	Mapping[str, Any] | None	Optional resolved configuration recorded before the first update.	None
ckpt_path	str | Path | None	Optional Orbax checkpoint root, or "last" for the configured callback's latest checkpoint.	None
ckpt_step	int | None	Exact checkpoint step, or None for the latest committed step.	None
weights_only	bool	Whether to restore only model weights into the fresh state.	False
datamodule	PhiDataModule | None	Optional DataModule supplying the source when batches is omitted.	None

Returns:

Type	Description
FitResult	Bound module, final state, metrics, and termination status.

Raises:

Type	Description
ValueError	If checkpoint restoration or DataModule source options are incomplete or invalid.
RuntimeError	If a custom training step does not increment TrainState.step exactly once.
BaseException	Re-raises training lifecycle errors after cleanup.

predict_state [source]

predict_state(
    module: BasePhiModule,
    state: TrainState,
    batches: Iterable[Mapping[str, Any]] | None = None,
    *,
    ckpt_path: str | Path | None = None,
    ckpt_step: int | None = None,
    metadata: Mapping[str, Any] | None = None,
    return_predictions: bool = True,
    datamodule: PhiDataModule | None = None,
) -> np.ndarray | None

Predict from an explicitly bound module and functional state template.

This follows Lightning's ckpt_path convention while preserving PhiJAX's explicit functional state. The supplied state defines the complete Orbax restore structure; only its model state is replaced for prediction. When batches is omitted, the trainer prepares and sets up datamodule for prediction and requests its batch source. Prediction is skipped when neither source is available.

Parameters:

Name	Type	Description	Default
module	BasePhiModule	Application module exposing prediction computation and lifecycle hooks.	required
state	TrainState	Fresh or in-memory functional state supplying the model and optional restore template.	required
batches	Iterable[Mapping[str, Any]] | None	Optional finite prediction batch iterable. When omitted, use datamodule.	None
ckpt_path	str | Path | None	Optional Orbax checkpoint root. None predicts from the supplied in-memory state.	None
ckpt_step	int | None	Exact checkpoint step, or None for the latest committed step.	None
metadata	Mapping[str, Any] | None	Optional immutable metadata exposed to prediction callbacks.	None
return_predictions	bool	Whether to concatenate and return host predictions.	True
datamodule	PhiDataModule | None	Optional DataModule set up when batches is omitted and torn down after prediction.	None

Returns:

Type	Description
ndarray | None	Concatenated host predictions, or None when collection is disabled or no prediction data is available.

Raises:

Type	Description
ValueError	If checkpoint options are invalid or a provided batch source is empty.
BaseException	Re-raises prediction lifecycle errors after cleanup.

initialize_state [source]

initialize_state(
    model_state: Any,
    optimizer: GradientTransformation,
    balancer_state: Any,
    key: Array,
    *,
    sampling_key: Array | None = None,
    balancer_key: Array | None = None,
) -> TrainState

Initialize functional state with this trainer's precision policy.

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 state.	required
key	Array	Root training key, or model-runtime key when both persistent keys are supplied.	required
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	Complete precision-aware training state.

compile_train_step [source]

compile_train_step(
    module: BasePhiModule,
    balancer: Any,
    optimizer: GradientTransformation,
) -> TrainStep

Build a compiled update with this trainer's precision policy.

Parameters:

Name	Type	Description	Default
module	BasePhiModule	Application module exposing named unweighted losses.	required
balancer	Any	Functional loss balancer.	required
optimizer	GradientTransformation	Optax gradient transformation used by initialized state.	required

Returns:

Type	Description
TrainStep	Reusable JIT-compiled training step.

prepare_batch_source [source]

prepare_batch_source(source: SourceT) -> SourceT

Prepare persistent sampler state for this Trainer's process-local root device.

Sources without a callable prepare method are returned unchanged, preserving support for ordinary callables and finite iterables. Built-in PhiJAX sources transfer finite candidate pools, generation bounds, templates, and PRNG keys once instead of transferring sampled training batches from the host on every step.

Parameters:

Name	Type	Description	Default
source	SourceT	Training batch callable or iterable, optionally exposing prepare(device).	required

Returns:

Type	Description
SourceT	Source prepared for the selected Strategy, or the original source when preparation is unsupported.

prepare_batch [source]

prepare_batch(batch: Any) -> Any

Cast and place one batch according to the configured precision and Strategy.

NumPy floating leaves are cast on the host before transfer. Device-resident sampler outputs retain their placement when the Strategy target is unchanged, while data-parallel Strategies apply final batch sharding.

Parameters:

Name	Type	Description	Default
batch	Any	Arbitrary host or device batch PyTree.	required

Returns:

Type	Description
Any	Precision-normalized batch placed or sharded for compiled execution.

load_weights [source]

load_weights(
    state: TrainState, step: int | None = None
) -> TrainState

Load model weights while preserving fresh optimizer and run state.

Parameters:

Name	Type	Description	Default
state	TrainState	Fresh target training state.	required
step	int | None	Checkpoint step, or None for the latest checkpoint.	None

Returns:

Type	Description
TrainState	Fresh state with restored model weights.

Raises:

Type	Description
ValueError	If checkpointing is not configured.

latest_checkpoint [source]

latest_checkpoint() -> tuple[Path, int]

Resolve the latest committed filesystem checkpoint.

Returns:

Type	Description
tuple[Path, int]	Checkpoint-root path and latest committed step.

Raises:

Type	Description
FileNotFoundError	If the checkpoint backend contains no committed state.
ValueError	If checkpointing is disabled or its backend has no filesystem directory.

callback_state_dict [source]

callback_state_dict() -> dict[str, Mapping[str, Any]]

Collect JSON-compatible persistent state from every callback.

Returns:

Type	Description
dict[str, Mapping[str, Any]]	Stable callback identifiers mapped to persistent callback state.

Raises:

Type	Description
TypeError	If a callback returns a non-mapping or non-JSON-compatible state.

load_callback_state_dict [source]

load_callback_state_dict(
    states: Mapping[str, Mapping[str, Any]],
) -> None

Restore callback states using stable identifiers.

Parameters:

Name	Type	Description	Default
states	Mapping[str, Mapping[str, Any]]	Checkpoint callback states.	required

Raises:

Type	Description
ValueError	If checkpoint and Trainer callback identifiers differ.

print_environment_info [source]

print_environment_info() -> None

Print Lightning-style precision and accelerator information on global rank zero.

close [source]

close() -> None

Release any checkpoint resources opened outside a Trainer task.

Fit and prediction stages already close their resources automatically. Calling this method is optional and remains supported for explicit lifecycle management and context-manager use.