Monitoring callbacks

Monitoring callbacks inspect completed training metrics on the Python host. They do not run inside the compiled update.

Early stopping

mode accepts "min" or "max", represented by the public MonitorMode alias. Use "min" for losses and "max" for scores where larger values are better.

EarlyStopping [source]

EarlyStopping(
    monitor: str = "train/loss",
    patience: int = 100,
    mode: MonitorMode = "min",
    min_delta: float = 0.0,
)

Bases: Callback

Stop training when a monitored scalar metric ceases to improve.

Initialize metric-based early stopping.

Parameters:

Name	Type	Description	Default
`monitor`	`str`	Metric name inspected after each training batch.	`'train/loss'`
`patience`	`int`	Number of non-improving observations tolerated.	`100`
`mode`	`MonitorMode`	Either `min` or `max`.	`'min'`
`min_delta`	`float`	Required absolute improvement over the current best.	`0.0`

Raises:

Type	Description
`ValueError`	If `patience`, `mode`, or `min_delta` is invalid.

on_train_batch_end [source]

on_train_batch_end(context: TrainerContext) -> bool

Update the monitored best value and request stopping after patience expires.

Parameters:

Name	Type	Description	Default
`context`	`TrainerContext`	Trainer context containing the latest metric mapping.	required

Returns:

Type	Description
`bool`	Whether the configured patience has expired.

Raises:

Type	Description
`KeyError`	If the monitored metric is absent.
`ValueError`	If the monitored metric is not scalar.

Learning-rate monitoring

LoggingInterval represents the explicit "step" and "epoch" interval values. Passing None follows the Trainer logging cadence.

LearningRateMonitor evaluates the Optax schedule at the optimizer count for each completed update. It sends optimizer/lr-Adam to configured loggers when optimizer_name="Adam". logging_interval="step" evaluates every step, while logging_interval="epoch" evaluates once at the end of fitting. The default None follows trainer.log_every_n_steps and always records the final rate. Unlike a Lightning scheduler, a raw Optax schedule has no interval field. PhiJAX therefore avoids evaluating the schedule on steps that will not be logged.

Like Lightning, this callback requires an experiment logger. Fitting raises before requesting the first batch when LearningRateMonitor is enabled with logger=False, logger=None, or an empty logger collection. The default logger=True is valid.

Optimizer names identify learning-rate series, while log_momentum and log_weight_decay add -momentum and -weight_decay suffixes. PhiJAX requires optimizer_name because Optax transformations do not retain their factory name. Metrics use the optimizer/ group by default; set log_key_prefix to replace it or pass None to disable grouping. Optional momentum and weight-decay values must be supplied because Optax transformations do not expose inspectable parameter groups.

LearningRateMonitor [source]

LearningRateMonitor(
    schedule: Callable[[int], Any] | None,
    *,
    optimizer_name: str,
    log_momentum: bool = False,
    log_weight_decay: bool = False,
    log_key_prefix: str | None = "optimizer/",
    logging_interval: LoggingInterval | None = None,
    momentum: Any = 0.0,
    weight_decay: Any = 0.0,
)

Bases: Callback

Publish optimizer hyperparameters through the standard trainer metric stream.

Optax schedules consume the pre-update optimizer count. A completed trainer step n therefore reports the value evaluated at n - 1, matching the hyperparameters used for that update. Unlike stateful PyTorch optimizers, an Optax transformation does not expose mutable parameter groups, so optional momentum and weight decay values are supplied explicitly by configuration.

Initialize optimizer hyperparameter monitoring.

Parameters:

Name	Type	Description	Default
`schedule`	`Callable[[int], Any] \| None`	Optax-compatible callable receiving a zero-based optimizer step. `None` supports bootstrap Hydra composition but is rejected when the callback is used for fitting.	required
`optimizer_name`	`str`	Non-empty display name added to `lr-<optimizer_name>`. PhiJAX requires this because Optax transformations do not retain an inspectable optimizer class name.	required
`log_momentum`	`bool`	Whether to log `momentum` as `lr-<optimizer_name>-momentum`.	`False`
`log_weight_decay`	`bool`	Whether to log `weight_decay` as `lr-<optimizer_name>-weight_decay`.	`False`
`log_key_prefix`	`str \| None`	Optional string prepended verbatim to every metric name. The default groups metrics under `optimizer/`; pass `None` to disable grouping.	`'optimizer/'`
`logging_interval`	`LoggingInterval \| None`	`"step"` evaluates every optimizer step, `"epoch"` evaluates at fit end, and `None` follows `trainer.log_every_n_steps` while always evaluating the final step.	`None`
`momentum`	`Any`	Optimizer momentum scalar or step-indexed schedule.	`0.0`
`weight_decay`	`Any`	Optimizer weight-decay scalar or step-indexed schedule.	`0.0`

Raises:

Type	Description
`TypeError`	If the schedule, logging flags, prefix, or optimizer hyperparameters have invalid types.
`ValueError`	If `optimizer_name` is empty or `logging_interval` is unsupported.

setup [source]

setup() -> None

Verify an executable fit supplied a learning-rate schedule.

Raises:

Type	Description
`RuntimeError`	If the callback was instantiated without a schedule.

on_fit_start [source]

on_fit_start(context: TrainerContext) -> None

Require an experiment logger before the first optimizer update.

Parameters:

Name	Type	Description	Default
`context`	`TrainerContext`	Initial Trainer context exposing logger availability.	required

Raises:

Type	Description
`RuntimeError`	If the Trainer has no configured experiment logger.

training_metrics [source]

training_metrics(
    context: TrainerContext,
) -> Mapping[str, jax.Array]

Evaluate the learning rate used by the completed optimizer update.

Parameters:

Name	Type	Description	Default
`context`	`TrainerContext`	Post-module context whose step counts completed optimizer updates.	required

Returns:

Type	Description
`Mapping[str, Array]`	Learning-rate metric and any enabled optimizer hyperparameter metrics.

Raises:

Type	Description
`RuntimeError`	If no schedule was configured.
`ValueError`	If a configured schedule does not return a scalar.

state_dict [source]

state_dict() -> Mapping[str, Any]

Return the most recently emitted optimizer step.

Returns:

Type	Description
`Mapping[str, Any]`	JSON-compatible monitor state.

load_state_dict [source]

load_state_dict(state: Mapping[str, Any]) -> None

Restore learning-rate emission bookkeeping.

Parameters:

Name	Type	Description	Default
`state`	`Mapping[str, Any]`	Mapping containing an integer or null `last_step`.	required

Raises:

Type	Description
`ValueError`	If the mapping shape or saved step is invalid.