Skip to content

Training

tinax.training owns one explicit training loop. It composes tinax.checkpointing with caller-owned step functions, batch iterables, evaluation callables, and reporting hooks; it is not a framework, and every policy — checkpoint cadence, evaluation cadence, metric monitoring, reporting — is an argument.

from tinax.training import fit

result = fit(
    mesh=mesh,
    model=model,
    optimizer=optimizer,
    step_fn=step,
    epoch_batches=batches,
    steps_per_epoch=32,
    epochs=10,
    checkpoint_dir="runs/exp-1",
    evaluate_fn=evaluate,
    monitor="accuracy",
    monitor_mode="max",
)

Checkpoints are written to checkpoint_dir/epoch-<NNNNNN> and never overwritten, so a later call resumes from the highest completed epoch. Batches must be deterministic per epoch, and only process 0 reports metrics.

fit 

fit(
    *,
    mesh: Mesh,
    model: Module,
    optimizer: Optimizer,
    step_fn: Callable[..., Mapping[str, Array]],
    epoch_batches: Callable[[int], Iterable[Any]],
    steps_per_epoch: int,
    epochs: int,
    checkpoint_dir: str | PathLike[str] | None = None,
    checkpoint_every: int = 1,
    evaluate_fn: Callable[[], Mapping[str, float]]
    | None = None,
    eval_every: int = 1,
    monitor: str | None = None,
    monitor_mode: Literal["min", "max"] = "max",
    log_metrics: Callable[[Mapping[str, float], int], None]
    | None = None,
    logger: Logger | None = None,
) -> FitResult

Train model for epochs epochs, evaluating, checkpointing, and reporting per epoch.

step_fn(model, optimizer, batch) performs one optimizer step and returns a mapping of scalar training metrics. epoch_batches(epoch) yields exactly steps_per_epoch batches for the zero-based epoch and must be deterministic so an epoch-level resume is exact. Checkpoints are written to checkpoint_dir/epoch-<NNNNNN> and never overwritten, so a later call resumes from the highest completed epoch. Only process 0 reports metrics.

Parameters:

Name Type Description Default
mesh Mesh

Device mesh the run is placed on.

 required
model Module

Flax NNX module updated in place by step_fn.

 required
optimizer Optimizer

Flax NNX optimizer updated in place by step_fn.

 required
step_fn Callable[..., Mapping[str, Array]]

Callable performing one optimizer step and returning scalar metrics.

 required
epoch_batches Callable[[int], Iterable[Any]]

Callable mapping a zero-based epoch to its batch iterable.

 required
steps_per_epoch int

Batches epoch_batches must yield for every epoch.

 required
epochs int

Total epochs to complete, counting any restored epochs.

 required
checkpoint_dir str | PathLike[str] | None

Directory holding epoch-<NNNNNN> checkpoints, or None to disable.

 None
checkpoint_every int

Write a checkpoint every N completed epochs.

 1
evaluate_fn Callable[[], Mapping[str, float]] | None

Zero-argument callable returning scalar evaluation metrics, or None.

 None
eval_every int

Evaluate every N completed epochs, plus the final epoch.

 1
monitor str | None

Evaluation metric name used to select the best epoch, or None.

 None
monitor_mode Literal['min', 'max']

Whether the monitored metric is maximized or minimized.

 'max'
log_metrics Callable[[Mapping[str, float], int], None] | None

Callable invoked once per epoch with (metrics, step), or None.

 None
logger Logger | None

Logger receiving one summary line per epoch, or None.

 None

Returns:

Type Description
FitResult

A FitResult with the final step, epoch, per-epoch history, and best epoch.

Raises:

Type Description
TypeError

If an argument has the wrong type, or step_fn/evaluate_fn return a non-mapping.
ValueError

If a count is not positive, monitor is set without evaluate_fn, monitor_mode is unknown, or an epoch yields the wrong number of batches.

FitResult dataclass 

FitResult(
    step: int,
    epoch: int,
    history: tuple[EpochRecord, ...],
    best: EpochRecord | None,
)

The outcome of one fit call.

Attributes:

Name Type Description
step int

Total optimizer steps completed, including restored ones.
epoch int

Total epochs completed, including restored ones.
history tuple[EpochRecord, ...]

One record per epoch completed by this call.
best EpochRecord | None

The record with the best monitored evaluation value, or None.

EpochRecord dataclass 

EpochRecord(
    epoch: int,
    step: int,
    train: Mapping[str, float],
    evaluation: Mapping[str, float] | None,
)

Metrics for one completed epoch.

Attributes:

Name Type Description
epoch int

One-based number of the completed epoch.
step int

Total optimizer steps completed after this epoch.
train Mapping[str, float]

Mean training metrics across the epoch.
evaluation Mapping[str, float] | None

Evaluation metrics, or None when the epoch was not evaluated.