LieselOptim

class liesel.optim.LieselOptim(model, *, loss_monitor, optimizers, stopper=None, seed=0, split=None, batch_size=None, batches=None, loss=None, validation_strategy='log_lik', scale_loss=True, save_position_history=True, show_progress=True, show_step_progress=False, progress_update_every=10, step_progress_update_every=10)

Bases: object

Builds an OptimEngine for a Liesel model using sensible defaults.

LieselOptim is the quick-start wrapper for regular model optimization. It creates a negative log-posterior loss, full-data training batches, and wraps a supplied Optax transformation over all model parameters. Call build_engine() to inspect or modify the low-level engine before fitting, or fit() for the direct path.

Parameters:

• model (Model) – Liesel model to optimize.

• 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 use the post-update epoch position. "train_full_data" always adds the full-data evaluation, including with one full-data optimization batch.

• optimizers (GradientTransformation | Sequence[OptimizerLike] | Literal['lbfgs']) – Required optimizer choice. An Optax transformation such as optax.adam(0.01) applies to all model parameters. Pass "lbfgs" for the built-in full-data L-BFGS optimizer, or a sequence of explicit optimizers for selected parameters. If a parameter is weak (computed), automatic selection raises an error. Explicitly name its strong source variables with Optimizer or LBFGS; those sources need not be marked as parameters. Priors on weak parameters remain in the loss. Pass a configured transformation, not an optimizer factory. Transformations must support updates from gradients, state and parameters without extra objective arguments.

• stopper (Stopper | None, default: None) – Maximum-epoch and early-stopping configuration. None creates a new Stopper with epochs=1000, patience=10, and rtol=1e-6.

• seed (int | None, default: 0) – Integer seed for batch shuffling or random batch sampling, and for custom losses or optimizers that use carry.key. Starting parameter values come from model. Set the seed for a random training/validation/test split separately when creating that split. Defaults to 0. Explicit None uses the current Unix time in whole seconds.

• split (PositionSplit | PositionSplitManager | None, default: None) – Optional split. If neither split nor loss is supplied, all observed data is used for training; weak observations are recomputed from their strong inputs. Observed arrays with different axis lengths require an explicit split. Check their row alignment before opting in with split=PositionSplit.from_model(model, multi_size="manager"). Use nested position_keys to specify independent groups and split_axes={key: None} for shared values. With a custom loss, the split must be the same object as loss.split. Models with per_obs=False require an explicitly constructed split: use PositionSplit.from_model with infer_sample_sizes=False for axis counts, or supply effective sample_sizes there. Custom aggregate likelihood or probability nodes require a custom loss as well as an explicit split.

• batch_size (int | None, default: None) – Rows per batch in each training group. Uses Batches.from_split with otherwise default settings. None uses all training data. Pass batches instead for custom settings, such as non-leading axes.

• batches (Batches | BatchManager | None, default: None) – Explicit batch configuration. Cannot be combined with a non-None batch_size.

• loss (Loss | None, default: None) – Custom loss. Uses loss.split and overrides validation_strategy and scale_loss.

• validation_strategy (Literal['log_lik', 'log_prob'], default: 'log_lik') – Validation strategy passed to NegLogProbLoss when loss is not supplied.

• scale_loss (bool, default: True) – Whether the default NegLogProbLoss should divide losses by the training sample size. Defaults to True. This setting has no effect when loss is supplied.

• save_position_history (bool, default: True) – Whether to save parameter values at every epoch. Defaults to True. History is allocated for the maximum epoch budget before fitting: its size is approximately epochs times the total parameter bytes (4 GB for 1,000 epochs and one million float32 parameters). Set False to skip this allocation; losses and final/best positions remain available.

• show_progress (bool, default: True) – Whether the built engine should show tqdm progress bars.

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

• progress_update_every (int, default: 10) – Update the epoch progress bar after this many completed epochs. When batch progress is active, the epoch bar advances after every epoch.

• step_progress_update_every (int, default: 10) – Update the batch progress bar after this many completed batches.

Examples

>>> import jax.numpy as jnp
>>> import optax
>>> import liesel.model as lsl
>>> import tensorflow_probability.substrates.jax.distributions as tfd
>>> from liesel.optim import EmaTrainLossMonitor, LieselOptim
>>> loc = lsl.Var.new_param(jnp.array(0.0), name="loc")
>>> y = lsl.Var.new_obs(
...     jnp.array([0.0, 1.0]),
...     lsl.Dist(tfd.Normal, loc=loc, scale=1.0),
...     name="y",
... )
>>> model = lsl.Model([y])
>>> engine = LieselOptim(
...     model,
...     optimizers=optax.adam(0.01),
...     loss_monitor=EmaTrainLossMonitor(effective_window=2.0),
...     seed=1,
... ).build_engine()
>>> type(engine).__name__
'OptimEngine'

Methods

build_engine()	Builds the low-level optimization engine.
fit()	Builds an engine and runs optimization immediately.

Attributes

batches	Explicit batch configuration.
loss	Custom loss.
loss_monitor	Source for the epoch-level stopping and progress loss.
model	Liesel model to optimize.
optimizers	Required optimizer choice.
progress_update_every	Update the epoch progress bar after this many completed epochs.
save_position_history	Whether to save parameter values at every epoch.
seed	Integer seed for batch shuffling or random batch sampling, and for custom losses or optimizers that use carry.key.
show_progress	Whether the built engine should 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	Optional split.
step_progress_update_every	Update the batch progress bar after this many completed batches.
stopper	Maximum-epoch and early-stopping configuration.