Batches

class liesel.optim.Batches(position_keys, axis_size, batch_size, shuffle=True, batch_axes=None, default_batch_axis=0, sample_with_replacement=False, sample_size=None, batch_sample_size=None, *, sampling_weights=None, likelihood_axes=None)

Bases: object

Defines mini-batches for observed entries in an optimizer position.

Batches owns the observation indices for the current epoch and reshapes complete parts of shuffled passes into batches. The observed position entries named in position_keys are sliced with these indices. By default, every entry is sliced along axis 0; use default_batch_axis or batch_axes for arrays where observations live on another axis.

Parameters:

• position_keys (Sequence[str]) – Names of the position entries that should be batched. An empty sequence is allowed only with batch_size=None; in that case the object is a full-data adapter that does not replace any observed entries.

• axis_size (int) – Number of observations along each batched axis.

• batch_size (int | None) – Number of observations per batch. If None, batching is disabled by using a single batch with all axis_size observations.

• shuffle (bool, default: True) – Whether permute_indices() should return a random permutation of the indices. If False, permute_indices() returns the indices unchanged.

• batch_axes (dict[str, int] | None, default: None) – Optional mapping from position key to batching axis. Keys missing from this mapping use default_batch_axis.

• default_batch_axis (int, default: 0) – Batching axis for all position keys not listed in batch_axes.

• sample_with_replacement (bool, default: False) – Whether every assembled batch draws observations independently with replacement. This is useful for explicit replacement sampling and is enabled automatically for an oversized child when a common batch size exceeds its observation count.

• sample_size (float | None, default: None) – Optional effective likelihood sample size represented by the full data. If omitted, likelihood scaling falls back to axis_size. Use from_model() to infer this value from pointwise observed log-probability arrays.

• batch_sample_size (float | None, default: None) – Optional effective likelihood sample size represented by one batch. If sample_size is supplied and batch_sample_size is omitted, it is derived as sample_size * batch_size / axis_size.

• sampling_weights (Any | None, default: None) – Optional finite, strictly positive relative weights of length axis_size, aligned with the training data in its current order. Requires sample_with_replacement=True. Normalized probabilities remain fixed during a run; saved probabilities and alias tables take precedence on checkpoint recovery. Probabilities and corrections use at least float32. Values whose probabilities or correction factors cannot be represented in the working dtype are rejected. No probability floor or clipping is applied.

• likelihood_axes (dict[str, int] | None, default: None) – Optional mapping from observed-variable name to its pointwise log-likelihood axis. Weighted correction normally infers this axis using trailing event reduction and leading broadcasting. Declare it explicitly for custom reductions or transpositions. The declared axis must enumerate the sampled observations; matching dimension sizes alone cannot establish this.

Notes

If axis_size is not divisible by batch_size, only full batches are used and the final incomplete tail is dropped independently from each shuffled pass. When more batches are requested, another independent shuffled pass supplies fresh rows; prior batches are never copied wholesale.

With replacement sampling, an epoch retains the same number of batches but does not guarantee coverage. Weighted groups apply an extra per-index factor 1 / (axis_size * p_i) through scaled_log_lik(), preserving the original objective in expectation. NegLogProbLoss applies this automatically. Other losses must use that method or apply correction_factors() before summing likelihood values.

A no-key full-data adapter can be useful when an optimizer workflow expects a Batches object but the model should always be evaluated on the full observed data. Use position_keys=[] and batch_size=None for this case.

Examples

Create two batches of size four from ten observations:

>>> from liesel.optim import Batches
>>> batches = Batches(["y"], axis_size=10, batch_size=4, shuffle=False)
>>> batches.batch_indices.tolist()
[[0, 1, 2, 3], [4, 5, 6, 7]]
>>> batches.n_full_batches
2

With batch_size=None, the object represents one full-data batch:

>>> full_data = Batches(["y"], axis_size=5, batch_size=None)
>>> full_data.batch_size
5
>>> full_data.batch_indices.tolist()
[[0, 1, 2, 3, 4]]

An empty-key full-data adapter leaves observed entries untouched:

>>> no_key = Batches([], axis_size=5, batch_size=None)
>>> no_key.position_keys, no_key.is_full_data
([], True)

batch_axes can batch different entries along different batch_axes:

>>> import jax.numpy as jnp
>>> batches = Batches(
...     ["x", "y"],
...     axis_size=5,
...     batch_size=2,
...     batch_axes={"x": 1},
...     shuffle=False,
... )
>>> position = {
...     "x": jnp.arange(15).reshape(3, 5),
...     "y": jnp.arange(20).reshape(5, 4),
... }
>>> batched = batches.get_batched_position(position, batch_index=0)
>>> batched["x"].shape, batched["y"].shape
((3, 2), (2, 4))

Methods

correction_factors(batch_index)	Return extra factors 1 / (axis_size * p_i) for the selected batch.
extract_batched_position(interface, ...)	Extracts observed data from a model state and returns one batch.
from_model(model, batch_size[, ...])
from_split(split, batch_size[, shuffle, ...])	Build batches from training data, preserving the split's groups.
get_batched_position(position, batch_index)	Slices observed position entries for one batch.
permute_indices(key)	Returns epoch indices, optionally shuffled.
scaled_log_lik(model, model_state, *[, ...])	Returns the log likelihood with this batch group's likelihood scaled.
start_epoch(key[, n_batches])	Starts a new epoch by updating the observation order.
weights_balanced(labels, *[, strength])	Construct sampling weights from category frequencies.
weights_binned(values, *, bins[, strength])	Construct sampling weights by balancing one-dimensional bin counts.
weights_for_shares(labels, shares, *[, ...])	Divide each category's target sampling share among its observations.

Attributes

batch_axes	Optional mapping from position key to batching axis.
batch_indices	Batch index matrix.
batch_sample_scale	Mini-batch likelihood scaling factor.
batch_sample_scales	Batch likelihood scaling factors.
batch_sample_size	Optional effective likelihood sample size represented by one batch.
default_batch_axis	Batching axis for all position keys not listed in batch_axes.
is_full_data	Whether the object represents one full-data batch.
n_full_batches	Number of complete batches.
sample_size	Optional effective likelihood sample size represented by the full data.
sample_with_replacement	Whether every assembled batch draws observations independently with replacement.
sampling_probabilities	Normalized sampling probabilities, or None for uniform sampling.
shuffle	Whether permute_indices() should return a random permutation of the indices.
position_keys	Names of the position entries that should be batched.
axis_size	Number of observations along each batched axis.
batch_size	Number of observations per batch.
likelihood_axes	Optional mapping from observed-variable name to its pointwise log-likelihood axis.
indices	Current observation order used by batch_indices; assign the result of permute_indices() to refresh it.