PositionSplit

class liesel.optim.PositionSplit(train, validate, test, train_axis_size, validate_axis_size, test_axis_size, sample_sizes=None, *, passthrough=<factory>)

Bases: object

Container for train, validation, and test position dictionaries.

A PositionSplit is usually returned by Split.split_position or PositionSplit.from_model. It stores the split observed position entries together with the corresponding split sizes.

Parameters:

• train (Position (dict[str, Any])) – Position entries used for optimization.

• validate (Position (dict[str, Any])) – Position entries used for validation or early stopping.

• test (Position (dict[str, Any])) – Position entries reserved for testing.

• train_axis_size (int) – Number of training observations.

• validate_axis_size (int) – Number of validation observations.

• test_axis_size (int) – Number of test observations.

• sample_sizes (Mapping[Literal['train', 'validate', 'test'], int | float] | None, default: None) – Optional effective sample sizes for train, validation, and test scaling. If omitted, split-axis counts are used. Use from_model() or add_inferred_sample_sizes_from_model() to infer these values from pointwise observed log-probability arrays.

• passthrough (Position (dict[str, Any]), default: <factory>) – Keyword-only position entries included unchanged in train, validate, and test. Use this for shared data that every split needs in full, such as a group-level lookup table: passthrough entries are not split, are excluded from split likelihood scaling, and are not batched automatically. Per-observation covariates, weights, and offsets should be split instead.

Examples

>>> import jax.numpy as jnp
>>> from liesel.optim import PositionSplit
>>> from liesel.optim.types import Position
>>> split = PositionSplit(
...     train=Position({"x": jnp.arange(3)}),
...     validate=Position({"x": jnp.arange(3, 5)}),
...     test=Position({"x": jnp.arange(5, 6)}),
...     train_axis_size=3,
...     validate_axis_size=2,
...     test_axis_size=1,
... )
>>> split
PositionSplit(train=3, validate=2, test=1)
>>> split.validate["x"].tolist()
[3, 4]

Methods

add_inferred_sample_sizes_from_model(model)	Infer effective sample sizes from model and attach them to this split.
from_model(model[, position_keys, ...])	Builds a PositionSplit from the observed variables in a model.
sample_scale(part)	Scale a part likelihood to the training sample size.
sample_size(part)	Effective likelihood sample size for one split part.
scaled_log_lik(model, model_state[, part])	Returns the log likelihood scaled for one split part.
set_sample_sizes(sample_sizes)	Set effective likelihood sample sizes on this split.

Attributes

axis_size	Total axis size represented by the split.
has_test	Whether this split contains test observations.
has_validation	Whether this split contains validation observations.
position_keys	Position keys contained in this split.
sample_sizes	Optional effective sample sizes for train, validation, and test scaling.
split_position_keys	Position keys that are partitioned rather than passed through.
test_axis_share	Share of observations assigned to testing.
test_sample_size	Effective test likelihood sample size.
train_sample_size	Effective training likelihood sample size.
validate_axis_share	Share of observations assigned to validation.
validate_sample_scale	Likelihood scale for validation data.
validate_sample_size	Effective validation likelihood sample size.
train	Position entries used for optimization.
validate	Position entries used for validation or early stopping.
test	Position entries reserved for testing.
train_axis_size	Number of training observations.
validate_axis_size	Number of validation observations.
test_axis_size	Number of test observations.
passthrough	Keyword-only position entries included unchanged in train, validate, and test.