OptimEngine

class liesel.optim.OptimEngine(loss, batches, optimizers, stopper, seed, initial_state, prune_history=True, show_progress=True, save_position_history=True, *, debug_nans=False, loss_monitor, progress_update_every=10, show_step_progress=False, step_progress_update_every=10)

Bases: object

Runs an optimization loop over epochs, batches, and optimizers.

OptimEngine runs the fit configured by LieselOptim. Each epoch starts by asking batches for fresh batch indices, then iterates over all full batches. For each batch, each optimizer gets a turn to update the subset of parameters named in its position_keys. The first active optimizer’s pre-update loss supplies the batch observation; if none is active, the loss is evaluated once at the unchanged position. At the end of the epoch, the engine records training and monitoring losses, updates the minimum-monitor position, and asks stopper whether to continue.

Parameters:

• loss (Loss) – Loss object implementing the Loss protocol.

• batches (Batches | BatchManager) – Batch configuration used for the training data. Use Batches for a single observation size and BatchManager for multi-branch models with different observation sizes.

• optimizers (Sequence[OptimizerLike]) – Sequence of optimizers. Each optimizer must claim a disjoint set of position keys. Individual optimizers may delay activation with Optimizer.activate_after_epochs.

• stopper (Stopper) – Early-stopping and maximum-epoch configuration.

• seed (int | Array) – Integer seed or JAX PRNG key for batch shuffling or random batch sampling, and for losses or optimizers that use carry.key. Starting parameters come from loss.position(); the data split is already defined. Resuming a checkpoint uses its saved random key.

• initial_state (Any) – Initial model state passed into OptimCarry.

• prune_history (bool, default: True) – If True, remove unused history entries after early stopping.

• show_progress (bool, default: True) – Whether to show tqdm progress bars. This is the master switch for both epoch and optional batch progress.

• save_position_history (bool, default: True) – Whether to store the full position history. The minimum-monitor position is tracked independently of this setting.

• loss_monitor (EmaTrainLossMonitor | Literal['validation', 'train_full_data']) – Source for the epoch-level stopping and progress loss. Pass EmaTrainLossMonitor for a continuous EMA of pre-update losses, "validation" for one complete validation-loss evaluation after each epoch, or "train_full_data" for one complete training-loss evaluation after each epoch. Exact monitors are evaluated at the final post-update epoch position, and "train_full_data" incurs this additional full-data evaluation even when optimization itself uses one full-data batch.

• debug_nans (bool, default: False) – Whether to capture first-NaN reproduction data during batch updates. Every active optimizer’s returned pre-update loss and the updated position are checked; when no optimizer is active, the explicit batched loss is checked.

• progress_update_every (int, default: 10) – Update the epoch progress bar after this many completed epochs. Defaults to 10. The final state is always rendered. When batch progress is active, the epoch bar advances after every epoch to keep the nested display consistent.

• show_step_progress (bool, default: False) – Whether to show an additional progress bar for batches within each epoch when show_progress is enabled.

• step_progress_update_every (int, default: 10) – Update the batch progress bar after this many completed batches. Defaults to 10. The final state of an interrupted epoch is always rendered.

Notes

OptimEngine uses carry.epoch as the number of completed epochs and as the next history index to be written. This matches Stopper’s indexing convention. Built-in LBFGS is accepted only with full-data batches and also requires a deterministic objective, which the engine cannot validate. Exact monitor minima retain the post-update position used for the evaluation. An EMA minimum instead retains the associated parameter snapshot; because an EMA combines losses from several positions, it is not an exact loss-position pairing.

Examples

Start with LieselOptim for ordinary fits. To assemble the pieces yourself:

>>> import jax.numpy as jnp
>>> import optax
>>> import tensorflow_probability.substrates.jax.distributions as tfd
>>> import liesel.model as lsl
>>> import liesel.optim as opt
>>> loc = lsl.Var.new_param(jnp.array(0.0), name="loc")
>>> y = lsl.Var.new_obs(
...     jnp.array([1.0, 2.0, 3.0]),
...     lsl.Dist(tfd.Normal, loc=loc, scale=1.0),
...     name="y",
... )
>>> model = lsl.Model([y])
>>> split = opt.PositionSplit.from_model(model)
>>> engine = opt.OptimEngine(
...     loss=opt.NegLogProbLoss(model, split, scale=True),
...     batches=opt.Batches.from_split(split, batch_size=None),
...     optimizers=[opt.Optimizer(["loc"], optax.adam(0.01))],
...     stopper=opt.Stopper(epochs=5, patience=5),
...     initial_state=model.state,
...     loss_monitor="train_full_data",
...     seed=42,
...     show_progress=False,
... )
>>> result = engine.fit()
>>> result.n_epochs
5

Methods

fit(*[, checkpoint, pause_after, ...])

Runs optimization and returns processed results.

Attributes

debug_nans	Whether to capture first-NaN reproduction data during batch updates.
position_keys	Position keys claimed by all optimizers.
progress_update_every	Update the epoch progress bar after this many completed epochs.
prune_history	If True, remove unused history entries after early stopping.
save_position_history	Whether to store the full position history.
show_progress	Whether to show tqdm progress bars.
show_step_progress	Whether to show an additional progress bar for batches within each epoch when show_progress is enabled.
split	Train/validation/test split supplied by loss.
step_progress_update_every	Update the batch progress bar after this many completed batches.
loss	Loss object implementing the Loss protocol.
batches	Batch configuration used for the training data.
optimizers	Sequence of optimizers.
stopper	Early-stopping and maximum-epoch configuration.
seed	Integer seed or JAX PRNG key for batch shuffling or random batch sampling, and for losses or optimizers that use carry.key.
initial_state	Initial model state passed into OptimCarry.
loss_monitor	Source for the epoch-level stopping and progress loss.