Display callbacks

The Trainer supplies plain terminal displays by default. Rich callbacks replace them when an application wants formatted tables and progress output.

Model summaries

Trainer(enable_model_summary=True) adds ModelSummary when no summary callback is supplied. Pass RichModelSummary to replace it, or set enable_model_summary=False for no automatic summary. A Trainer rejects multiple summary callbacks and prints summaries only on global rank zero.

ModelSummary [source]

ModelSummary(
    *,
    max_depth: int = -1,
    console_width: int = 120,
    compute_flops: bool = False,
    compute_vjp_flops: bool = False,
    rank_zero_only: bool = True,
    warn_if_unavailable: bool = True,
)

Bases: Callback

Print a network summary at fit start using the module's configured summary provider.

Initialize model-summary display policy.

Parameters:

Name	Type	Description	Default
`max_depth`	`int`	Maximum displayed module depth, or `-1` for every level.	`-1`
`console_width`	`int`	Positive Rich console width in characters.	`120`
`compute_flops`	`bool`	Whether to estimate forward-pass floating-point operations.	`False`
`compute_vjp_flops`	`bool`	Whether to estimate reverse-pass floating-point operations.	`False`
`rank_zero_only`	`bool`	Whether to suppress output outside global rank zero.	`True`
`warn_if_unavailable`	`bool`	Whether to warn when no summary can be rendered.	`True`

Raises:

Type	Description
`ValueError`	If `max_depth` is below `-1` or `console_width` is not positive.

on_fit_start [source]

on_fit_start(context: TrainerContext) -> None

Print the configured model summary before the first training batch.

Parameters:

Name	Type	Description	Default
`context`	`TrainerContext`	Initial trainer context containing the module and explicit model state.	required

RichModelSummary [source]

RichModelSummary(
    *,
    max_depth: int = -1,
    console_width: int = 120,
    compute_flops: bool = False,
    compute_vjp_flops: bool = False,
    rank_zero_only: bool = True,
    warn_if_unavailable: bool = True,
)

Bases: ModelSummary

Select the Rich-capable model summary provider as an explicit display callback.

Progress bars

When no progress callback is supplied, Trainer(enable_progress_bar=True) adds TQDMProgressBar. Set enable_progress_bar=False for quiet library, test, or batch-scheduler execution.

All progress callbacks provide enable(), disable(), and is_enabled for temporary runtime control. Most applications should use the Trainer option instead of managing callback state directly.

ProgressBar [source]

ProgressBar(
    total: int,
    *,
    refresh_rate: int = 10,
    metric_names: Sequence[str] | None = None,
    description: str = "Training",
    predict_description: str = "Predicting",
    rank_zero_only: bool = True,
    metrics_format: str = ".3e",
)

Bases: Callback

Define shared enablement and metric selection for progress displays.

Initialize common progress display policy.

Parameters:

Name	Type	Description	Default
`total`	`int`	Positive maximum optimizer updates in one fit call.	required
`refresh_rate`	`int`	Positive number of batches between metric refreshes.	`10`
`metric_names`	`Sequence[str] \| None`	Optional ordered exact metric names overriding module-selected display metrics.	`None`
`description`	`str`	Non-empty fit progress label.	`'Training'`
`predict_description`	`str`	Non-empty prediction progress label.	`'Predicting'`
`rank_zero_only`	`bool`	Whether only JAX process zero writes output.	`True`
`metrics_format`	`str`	Valid scalar format specification.	`'.3e'`

Raises:

Type	Description
`ValueError`	If a size, label, metric name, or scalar format is invalid.

is_enabled `property`

is_enabled: bool

Return whether rendering is enabled.

Returns:

Type	Description
`bool`	Whether the callback may render output.

is_disabled `property`

is_disabled: bool

Return whether rendering is disabled.

Returns:

Type	Description
`bool`	Whether the callback suppresses output.

connect [source]

connect(trainer: Any) -> None

Connect the progress callback to its Trainer.

Parameters:

Name	Type	Description	Default
`trainer`	`Any`	Trainer exposing routed progress metrics and logger version.	required

disable [source]

disable() -> None

Disable progress rendering.

enable [source]

enable() -> None

Enable progress rendering.

get_metrics [source]

get_metrics(
    trainer: Trainer, module: BasePhiModule | None
) -> dict[str, Any]

Combine standard Trainer fields with module-selected progress metrics.

Parameters:

Name	Type	Description	Default
`trainer`	`Trainer`	Owning Trainer.	required
`module`	`BasePhiModule \| None`	Bound module participating in training, when available.	required

Returns:

Type	Description
`dict[str, Any]`	Ordered host or device values selected for display.

TQDMProgressBar [source]

TQDMProgressBar(
    total: int,
    *,
    stream: TextIO | None = None,
    **kwargs: Any,
)

Bases: ProgressBar

Display compact Lightning-style fit and prediction progress with TQDM.

Initialize a TQDM progress callback.

Parameters:

Name	Type	Description	Default
`total`	`int`	Positive maximum optimizer updates in one fit call.	required
`stream`	`TextIO \| None`	Optional output stream, defaulting to the current `sys.stderr` during setup.	`None`
`**kwargs`	`Any`	Common progress options forwarded to :class:`ProgressBar`.	`{}`

setup [source]

setup() -> None

Reset task-local TQDM state and resolve rank-owned output.

on_fit_start [source]

on_fit_start(context: TrainerContext) -> None

Create the fit progress display.

Parameters:

Name	Type	Description	Default
`context`	`TrainerContext`	Initial Trainer context.	required

on_train_metrics [source]

on_train_metrics(context: TrainerContext) -> None

Advance training progress and periodically refresh selected metrics.

Parameters:

Name	Type	Description	Default
`context`	`TrainerContext`	Complete post-update metric context.	required

on_fit_end [source]

on_fit_end(context: TrainerContext) -> None

Render terminal fit metrics and close the bar.

Parameters:

Name	Type	Description	Default
`context`	`TrainerContext`	Terminal fit context.	required

on_predict_start [source]

on_predict_start(context: PredictionContext) -> None

Create the prediction progress display.

Parameters:

Name	Type	Description	Default
`context`	`PredictionContext`	Initial prediction context.	required

on_predict_batch_end [source]

on_predict_batch_end(context: PredictionContext) -> None

Advance prediction progress by one batch.

Parameters:

Name	Type	Description	Default
`context`	`PredictionContext`	Completed prediction context.	required

on_predict_end [source]

on_predict_end(context: PredictionContext) -> None

Close the completed prediction display.

Parameters:

Name	Type	Description	Default
`context`	`PredictionContext`	Terminal prediction context.	required

on_exception [source]

on_exception(
    exception: BaseException, context: CallbackContext
) -> None

Mark a failed task and close its display.

Parameters:

Name	Type	Description	Default
`exception`	`BaseException`	Exception terminating the task.	required
`context`	`CallbackContext`	Most recent callback context.	required

teardown [source]

teardown() -> None

Close TQDM resources idempotently.

Supplying RichProgressBar, which extends ProgressBar, replaces TQDM. A Trainer accepts at most one progress callback, avoiding duplicate output. Both implementations refresh device metrics only at their configured interval.

By default, both displays show train/loss, every train/loss/<name>, every train/weight/<name>, and the primary logger version as v_num. Learning rates and other diagnostics still reach loggers without entering the display. Override ProgressBar.get_metrics() to change standard fields, or pass metric_names for an exact ordered selection. During prediction, the callbacks use PredictionContext.total_batches without transferring prediction values to the host.

RichProgressBarTheme [source] `dataclass`

RichProgressBarTheme(
    description: str = "",
    progress_bar: str = "#6206E0",
    progress_bar_finished: str = "#6206E0",
    progress_bar_pulse: str = "#6206E0",
    batch_progress: str = "",
    time: str = "dim",
    processing_speed: str = "dim underline",
    metrics: str = "italic",
    metrics_text_delimiter: str = " ",
    metrics_format: str = ".3e",
)

Configure Lightning-style colors and formatting for :class:RichProgressBar.

Attributes:

Name	Type	Description
`description`	`str`	Rich style applied to the step description.
`progress_bar`	`str`	Rich style applied to the active portion of the bar.
`progress_bar_finished`	`str`	Rich style applied to a completed bar.
`progress_bar_pulse`	`str`	Rich style applied when the bar pulses.
`batch_progress`	`str`	Rich style applied to the completed-step counter.
`time`	`str`	Rich style applied to elapsed and remaining time.
`processing_speed`	`str`	Rich style applied to iteration throughput.
`metrics`	`str`	Rich style applied to displayed metrics.
`metrics_text_delimiter`	`str`	Text separating adjacent metrics.
`metrics_format`	`str`	Python format specification applied to scalar metrics.

RichProgressBar [source]

RichProgressBar(
    total: int,
    *,
    refresh_rate: int = 10,
    metric_names: Sequence[str] | None = None,
    description: str = "Training",
    predict_description: str = "Predicting",
    transient: bool = False,
    rank_zero_only: bool = True,
    theme: RichProgressBarTheme | None = None,
    console: Console | None = None,
)

Bases: ProgressBar

Display fit and prediction progress with Rich.

The batch counter advances without transferring device metrics to the host. Selected metric values are copied only on the first batch, every refresh_rate batches, and at fit completion. This keeps the progress display useful without introducing a device synchronization on every accelerator update.

Initialize the Rich training progress display.

Parameters:

Name	Type	Description	Default
`total`	`int`	Number of batches expected in one fit call.	required
`refresh_rate`	`int`	Positive interval between device-to-host metric transfers.	`10`
`metric_names`	`Sequence[str] \| None`	Optional ordered exact names to display. `None` uses module-selected progress metrics.	`None`
`description`	`str`	Non-empty label displayed before the progress bar.	`'Training'`
`predict_description`	`str`	Non-empty label displayed during prediction.	`'Predicting'`
`transient`	`bool`	Whether Rich removes the progress display after completion.	`False`
`rank_zero_only`	`bool`	Whether only JAX process zero renders the display.	`True`
`theme`	`RichProgressBarTheme \| None`	Optional Lightning-style colors and scalar formatting.	`None`
`console`	`Console \| None`	Optional Rich console, primarily useful for embedding and tests.	`None`

Raises:

Type	Description
`ValueError`	If `total`, `refresh_rate`, either description, or a metric name is invalid.

setup [source]

setup() -> None

Create fit-local Rich resources on the process responsible for display.

on_fit_start [source]

on_fit_start(context: TrainerContext) -> None

Start a new progress task before the first training batch.

Parameters:

Name	Type	Description	Default
`context`	`TrainerContext`	Initial trainer context. Its global step may be nonzero after resume.	required

on_train_metrics [source]

on_train_metrics(context: TrainerContext) -> None

Advance the bar and periodically refresh its selected metrics.

Parameters:

Name	Type	Description	Default
`context`	`TrainerContext`	Post-update trainer context containing device metrics.	required

on_fit_end [source]

on_fit_end(context: TrainerContext) -> None

Render terminal metrics and stop the progress display.

Parameters:

Name	Type	Description	Default
`context`	`TrainerContext`	Final trainer context after a successful or callback-stopped fit.	required

on_predict_start [source]

on_predict_start(context: PredictionContext) -> None

Start a progress task for one finite prediction source.

Parameters:

Name	Type	Description	Default
`context`	`PredictionContext`	Initial prediction context containing the batch count when known.	required

on_predict_batch_end [source]

on_predict_batch_end(context: PredictionContext) -> None

Advance the prediction bar after one valid output batch.

Parameters:

Name	Type	Description	Default
`context`	`PredictionContext`	Completed prediction-batch context.	required

on_predict_end [source]

on_predict_end(context: PredictionContext) -> None

Render the completed prediction count and stop the progress display.

Parameters:

Name	Type	Description	Default
`context`	`PredictionContext`	Final prediction context.	required

on_exception [source]

on_exception(
    exception: BaseException, context: CallbackContext
) -> None

Mark and stop the progress display before an exception is re-raised.

Parameters:

Name	Type	Description	Default
`exception`	`BaseException`	Exception raised during training.	required
`context`	`CallbackContext`	Most recent valid callback context.	required

teardown [source]

teardown() -> None

Stop Rich resources idempotently during trainer cleanup.

Display callbacks

Model summaries

ModelSummary [source]

on_fit_start [source]

RichModelSummary [source]

Progress bars

ProgressBar [source]

is_enabled property

is_disabled property

connect [source]

disable [source]

enable [source]

get_metrics [source]

TQDMProgressBar [source]

setup [source]

on_fit_start [source]

on_train_metrics [source]

on_fit_end [source]

on_predict_start [source]

on_predict_batch_end [source]

on_predict_end [source]

on_exception [source]

teardown [source]

RichProgressBarTheme [source] dataclass

RichProgressBar [source]

setup [source]

on_fit_start [source]

on_train_metrics [source]

on_fit_end [source]

on_predict_start [source]

on_predict_batch_end [source]

on_predict_end [source]

on_exception [source]

teardown [source]

is_enabled `property`

is_disabled `property`

RichProgressBarTheme [source] `dataclass`