ML Inference¶

Use rollout_batch when a compatible model and tensors are already available. Use evaluate_checkpoints to materialize selected normalized streams once and compare several checkpoints and rollout suffix widths.

The comparative entry point accepts rows from an ML index:

import polars as pl
import torch

from batgrad.contracts.mapping import BaseColumns
from batgrad.ml.config import load_experiment_config, resolve_store_root
from batgrad.ml.data.loader import create_index
from batgrad.ml.inference import CheckpointSelection, evaluate_checkpoints
from batgrad.storage.local import LocalDataProcessingStore

config = load_experiment_config("configs/ml_baseline.json")
store = LocalDataProcessingStore(resolve_store_root(config.data.store_root))
index = create_index(
    store,
    config.data.manifest_paths,
    protocols=config.data.protocols,
    protocol_mode=config.data.protocol_mode,
)
selected = index.frame.filter(pl.col(BaseColumns.proto) == "cycling").head(1)

result = evaluate_checkpoints(
    store,
    selected,
    (CheckpointSelection("baseline", "outputs/runs/run/checkpoints/run/final.pt"),),
    device=torch.device("cuda:0"),
    suffix_steps=(0, config.train.masked_suffix.suffix_steps),
    rollout_steps=256,
)
if result.warning:
    print(result.warning)

Suffix width 0 selects classic one-step rollout. Positive widths use masked suffix rollout and must be smaller than the checkpoint context length. Returned predictions are decoded in the configured output-scaled units; inverse-scale them before comparing with physical measurements. Comparative inference starts at source-row offset zero for every selected stream and does not currently support EIS.

rollout_batch ¶

rollout_batch(
    config: ExperimentConfig,
    model: Module,
    inputs: Tensor,
    *,
    context_len: int,
    rollout_steps: int,
    suffix: MaskedSuffixConfig,
    device: device,
    targets: Tensor | None = None,
    mask: Tensor | None = None,
) -> RolloutResult

Run classic or masked-suffix autoregressive rollout for a batch.

Known future non-feedback controls must already be present in inputs. Feedback columns are replaced recursively with detached model predictions. Supplying both targets and mask enables scored evaluation; omitting both performs deployment-style prediction with identical execution semantics.

Parameters:

config (ExperimentConfig) –

Experiment defining columns, scaling, encoding, and feedback.
model (Module) –

Compatible model in evaluation mode.
inputs (Tensor) –

Scaled scalar inputs shaped (B, T, C_in) containing context and enough future control rows for the requested horizon.
context_len (int) –

Number of initially observed input rows.
rollout_steps (int) –

Maximum future rows to predict.
suffix (MaskedSuffixConfig) –

Disabled for one-step rollout or enabled for chunked suffix calls.
device (device) –

Model execution device.
targets (Tensor | None, default: None ) –

Optional scaled next-row targets shaped (B, T, C_out).
mask (Tensor | None, default: None ) –

Optional valid-target mask supplied together with targets.

Returns:

RolloutResult –

Decoded output-scaled predictions, optional metrics, and target alignment.

Raises:

ValueError –

If shapes, context, horizon, suffix width, or optional scoring inputs are invalid.

Examples:

Run deployment-style one-step feedback without metrics:

from batgrad.ml.config import MaskedSuffixConfig
from batgrad.ml.rollout import rollout_batch

result = rollout_batch(
    config,
    model,
    inputs,
    context_len=512,
    rollout_steps=64,
    suffix=MaskedSuffixConfig(enabled=False),
    device=device,
)
assert result.prediction.shape[1] == 64

Note

Attention recomputes over each visible context. Mamba state is advanced by exactly the rows leaving the next window and must not be shared across unrelated streams.

RolloutResult `dataclass` ¶

RolloutResult(
    prediction: Tensor,
    metrics: LossMetrics | None,
    target_start: int,
)

Predictions and optional metrics from one batched rollout.

Attributes:

prediction (Tensor) –

Decoded output-scaled values shaped (B, executed_steps, C_out).
metrics (LossMetrics | None) –

Count-weighted metrics when targets and a mask were supplied.
target_start (int) –

Target-sequence offset aligned with prediction step zero; equal to context_len - 1 under the next-row target contract.

evaluate_checkpoints ¶

evaluate_checkpoints(
    store: DatasetStoreReader,
    selected_index_frame: DataFrame,
    selections: tuple[CheckpointSelection, ...],
    *,
    device: device,
    suffix_steps: tuple[int, ...],
    rollout_steps: int,
) -> InferenceResult

Evaluate compatible checkpoints on selected normalized streams.

The selected index rows are materialized once, then every checkpoint is run with every requested suffix width. Checkpoints must agree on model-relevant columns, sequence length, and scaling.

Parameters:

store (DatasetStoreReader) –

Reader for normalized stream data.
selected_index_frame (DataFrame) –

One or more rows from an MlDatasetIndex.
selections (tuple[CheckpointSelection, ...]) –

Named checkpoint paths.
device (device) –

Evaluation device.
suffix_steps (tuple[int, ...]) –

Suffix widths to compare; zero requests classic rollout.
rollout_steps (int) –

Requested prediction horizon.

Returns:

InferenceResult –

CPU context/targets and comparative predictions. The horizon is clipped to
InferenceResult –

the shortest selected stream and reported through warning.

Raises:

ValueError –

If selections are empty, incompatible, unsupported by selected protocols, or too short for the checkpoint context.

Note

EIS inference is not currently supported. All selected rows must match a protocol configured by every checkpoint. Materialization always starts at source-row offset zero for each selected stream.

resolve_device ¶

resolve_device(value: str) -> device

Validate and return a requested PyTorch execution device.

Parameters:

value (str) –

Device string such as "cpu" or "cuda:0".

Returns:

device –

The parsed PyTorch device.

Raises:

ValueError –

If a requested CUDA device is unavailable.

available_devices ¶

available_devices() -> tuple[str, ...]

Return CPU and currently available indexed CUDA device strings.

Returns:

tuple[str, ...] –

A tuple beginning with "cpu", followed by available "cuda:N" devices.

CheckpointSelection `dataclass` ¶

CheckpointSelection(alias: str, path: str)

Named checkpoint included in a comparative inference request.

Attributes:

alias (str) –

Non-empty display label for plots and metrics.
path (str) –

Checkpoint path.

InferencePrediction `dataclass` ¶

InferencePrediction(
    checkpoint_alias: str,
    checkpoint_path: str,
    suffix_steps: int,
    context_predictions: Tensor,
    predictions: Tensor,
    metrics: LossMetrics | None,
    target_start: int,
)

One checkpoint and suffix-setting prediction series.

Attributes:

checkpoint_alias (str) –

User-provided checkpoint label.
checkpoint_path (str) –

Loaded checkpoint path.
suffix_steps (int) –

Requested suffix width; zero denotes classic one-step rollout.
context_predictions (Tensor) –

CPU context predictions shaped (B, context_len, C_out).
predictions (Tensor) –

CPU tensor shaped (B, rollout_len, C_out).
metrics (LossMetrics | None) –

CPU count-weighted metrics for observed targets.
target_start (int) –

Target offset aligned with prediction step zero.

InferenceResult `dataclass` ¶

InferenceResult(
    config: ExperimentConfig,
    inputs: Tensor,
    targets: Tensor,
    predictions: tuple[InferencePrediction, ...],
    context_len: int,
    rollout_len: int,
    group_keys: tuple[tuple[object, ...], ...],
    warning: str | None,
)

Materialized context and comparative checkpoint predictions.

Attributes:

config (ExperimentConfig) –

Shared compatible experiment configuration.
inputs (Tensor) –

Materialized scaled inputs on the CPU.
targets (Tensor) –

Materialized scaled next-row targets on the CPU.
predictions (tuple[InferencePrediction, ...]) –

Results for each checkpoint and requested suffix width.
context_len (int) –

Checkpoint sequence length.
rollout_len (int) –

Effective horizon after shortest-stream clipping.
group_keys (tuple[tuple[object, ...], ...]) –

Stream keys aligned with the batch dimension.
warning (str | None) –

Human-readable clipping warning, if any.

Checkpoints¶

discover_checkpoints ¶

discover_checkpoints(
    root: str | Path = ".",
) -> tuple[Path, ...]

Discover training checkpoints below a root directory.

The search includes both direct checkpoints/*.pt files and the checkpoints/<run-id-or-name>/*.pt layout produced by training.

Parameters:

root (str | Path, default: '.' ) –

Directory containing one or more run directories.

Returns:

tuple[Path, ...] –

Sorted checkpoint paths.

read_checkpoint_config ¶

read_checkpoint_config(
    path: str | Path,
) -> ExperimentConfig

Read and validate only the experiment configuration from a checkpoint.

Parameters:

path (str | Path) –

PyTorch checkpoint path.

Returns:

ExperimentConfig –

The embedded experiment configuration.

Note

PyTorch still deserializes the checkpoint payload; this helper merely avoids constructing the model.

load_checkpoint ¶

load_checkpoint(
    path: str | Path, device: device
) -> LoadedCheckpoint

Reconstruct a trained model for evaluation.

Parameters:

path (str | Path) –

PyTorch checkpoint path.
device (device) –

Device used for deserialization and model construction.

Returns:

LoadedCheckpoint –

The validated configuration, evaluation-mode model, and optional step.

Raises:

TypeError –

If required configuration or model payloads are absent.
RuntimeError –

If model weights are incompatible with the configuration.

Note

Checkpoints are loaded through PyTorch's restricted weights-only unpickler.

LoadedCheckpoint `dataclass` ¶

LoadedCheckpoint(
    config: ExperimentConfig,
    model: Module,
    step: int | None,
)

Model and experiment metadata restored for evaluation.

Attributes:

config (ExperimentConfig) –

Validated configuration embedded in the checkpoint.
model (Module) –

Reconstructed model in evaluation mode on the requested device.
step (int | None) –

Saved training step, if present and valid.

Metrics¶

LossMetrics `dataclass` ¶

LossMetrics(
    loss: Tensor,
    feature_loss_sum: Tensor | None = None,
    feature_loss_count: Tensor | None = None,
    feature_squared_error_sum: Tensor | None = None,
    feature_squared_error_count: Tensor | None = None,
    mamba_states: dict[str, MambaCarryState] | None = None,
)

Count-weighted categorical loss and decoded-error components.

Sums and counts are retained instead of batch means so callers can aggregate batches, rollout chunks, and distributed ranks without mean-of-means bias.

Attributes:

loss (Tensor) –

Scalar categorical cross-entropy over all valid selected values.
feature_loss_sum (Tensor | None) –

Per-target cross-entropy sums.
feature_loss_count (Tensor | None) –

Per-target valid-value counts.
feature_squared_error_sum (Tensor | None) –

Per-target decoded squared-error sums in output-scaled units.
feature_squared_error_count (Tensor | None) –

Per-target decoded-error counts.
mamba_states (dict[str, MambaCarryState] | None) –

Optional recurrent states produced by objective execution.

ML Inference¶

rollout_batch ¶

RolloutResult dataclass ¶

evaluate_checkpoints ¶

resolve_device ¶

available_devices ¶

CheckpointSelection dataclass ¶

InferencePrediction dataclass ¶

InferenceResult dataclass ¶

Checkpoints¶

discover_checkpoints ¶

read_checkpoint_config ¶

load_checkpoint ¶

LoadedCheckpoint dataclass ¶

Metrics¶

LossMetrics dataclass ¶

RolloutResult `dataclass` ¶

CheckpointSelection `dataclass` ¶

InferencePrediction `dataclass` ¶

InferenceResult `dataclass` ¶

LoadedCheckpoint `dataclass` ¶

LossMetrics `dataclass` ¶