ML Models¶

The built-in sequence mixer represents scalar values as categorical distributions, projects each input feature independently, reduces the feature axis, and applies configured temporal layers. The main path must contain exactly one reduction layer at index zero, and the head cannot contain reduction layers. Mamba layers require Linux, CUDA, and the ml dependency group.

SequenceMixerConfig `dataclass` ¶

SequenceMixerConfig(
    d_model: int = 256,
    n_heads: int = 8,
    mlp_ratio: float = 4.0,
    dropout: float = 0.0,
    bias: bool = False,
    norm: Literal["rmsnorm"] = "rmsnorm",
    causal_attention: bool = True,
    num_bins: int = 64,
    input_sigma: float = 0.0,
    output_sigma: float = 0.0,
    feedback_mode: Literal[
        "probabilities", "decoded_scalar"
    ] = "probabilities",
    mamba: MambaConfig = MambaConfig(),
    layers: tuple[LayerConfig, ...] = (
        LayerConfig(kind="reduce", mode="sum_pool"),
        LayerConfig(kind="attention"),
        LayerConfig(kind="ffn"),
    ),
    head_layers: tuple[LayerConfig, ...] = (
        LayerConfig(kind="ffn"),
    ),
    output: OutputConfig = OutputConfig(),
)

Feature encoding and sequence-mixer architecture.

Attributes:

d_model (int) –

Temporal token width.
n_heads (int) –

Attention head count; must divide d_model.
mlp_ratio (float) –

Feed-forward hidden-width multiplier.
dropout (float) –

Attention and feed-forward dropout probability.
bias (bool) –

Default bias setting for Batgrad-owned linear projections.
norm (Literal['rmsnorm']) –

Normalization implementation; currently RMSNorm only.
causal_attention (bool) –

Prevent attention to future sequence positions.
num_bins (int) –

Categorical bins used to encode each scalar.
input_sigma (float) –

Gaussian width for input encoding; zero linearly interpolates between adjacent bins.
output_sigma (float) –

Gaussian width for target distributions.
feedback_mode (Literal['probabilities', 'decoded_scalar']) –

Reuse detached probabilities directly or decode and re-encode scalar feedback.
mamba (MambaConfig) –

Default Mamba parameters.
layers (tuple[LayerConfig, ...]) –

Main feature-reduction and temporal path.
head_layers (tuple[LayerConfig, ...]) –

Temporal layers after feature reduction.
output (OutputConfig) –

Categorical output projection settings.

Examples:

Build a small causal attention model:

from batgrad.ml.nn import LayerConfig, SequenceMixerConfig

config = SequenceMixerConfig(
    d_model=128,
    n_heads=4,
    layers=(
        LayerConfig(kind="reduce", mode="sum_pool"),
        LayerConfig(kind="attention"),
        LayerConfig(kind="ffn"),
    ),
)

LayerConfig `dataclass` ¶

LayerConfig(
    kind: LayerKind,
    residual: bool | ResidualConfig | None = None,
    bias: bool | None = None,
    mode: Literal["sum_pool"] | None = None,
    d_state: int | None = None,
    expand: int | None = None,
    headdim: int | None = None,
    ngroups: int | None = None,
    is_mimo: bool | None = None,
    mimo_rank: int | None = None,
    chunk_size: int | None = None,
)

One feature-reduction or temporal-mixing layer.

Attributes:

kind (LayerKind) –

Attention, feed-forward, Mamba, or feature reduction.
residual (bool | ResidualConfig | None) –

Residual override. Temporal layers default to enabled and reduction defaults to disabled.
bias (bool | None) –

Optional attention/FFN linear-bias override. None inherits the model-level setting.
mode (Literal['sum_pool'] | None) –

Reduction implementation; currently only "sum_pool".
d_state (int | None) –

Optional per-layer Mamba state dimension.
expand (int | None) –

Optional per-layer Mamba expansion factor.
headdim (int | None) –

Optional per-layer Mamba head width.
ngroups (int | None) –

Optional per-layer Mamba group count.
is_mimo (bool | None) –

Optional per-layer Mamba MIMO mode.
mimo_rank (int | None) –

Optional per-layer Mamba projection rank.
chunk_size (int | None) –

Optional per-layer Mamba kernel chunk size.

Note

The main layer sequence must contain exactly one reduction layer at index zero. Head layers cannot reduce. This validated transition changes (B, T, C_in, D) to (B, T, D) before temporal layers.

ResidualConfig `dataclass` ¶

ResidualConfig(kind: ResidualKind = 'standard')

Residual connection policy for a configured layer.

Attributes:

kind (ResidualKind) –

Standard additive residual or no residual.

MambaConfig `dataclass` ¶

MambaConfig(
    d_state: int = 128,
    expand: int = 2,
    headdim: int = 64,
    ngroups: int = 1,
    is_mimo: bool = False,
    mimo_rank: int = 1,
    chunk_size: int = 64,
)

Default Mamba-3 layer parameters.

Attributes:

d_state (int) –

State-space dimension.
expand (int) –

Inner expansion factor.
headdim (int) –

Per-head width.
ngroups (int) –

Number of state-space parameter groups.
is_mimo (bool) –

Enable Mamba-3 MIMO mode.
mimo_rank (int) –

MIMO projection rank.
chunk_size (int) –

Kernel processing chunk size.

Note

Mamba layers require Linux, CUDA, and the ml dependency group. Cross-window state carry currently requires SISO mode (is_mimo=False).

OutputConfig `dataclass` ¶

OutputConfig(
    parameterization: Literal["shared"] = "shared",
)

Categorical output-head configuration.

Attributes:

parameterization (Literal['shared']) –

Output parameter sharing mode. Only "shared" is currently supported.

build_model ¶

build_model(
    config: ExperimentConfig, device: device
) -> SequenceMixer

Build an experiment's model on the requested device.

Parameters:

config (ExperimentConfig) –

Validated experiment configuration defining architecture and ordered input/target columns.
device (device) –

Destination device.

Returns:

SequenceMixer –

An initialized sequence mixer on device.

Raises:

ValueError –

If a Mamba model is requested on a non-CUDA device.

SequenceMixer ¶

SequenceMixer(
    config: SequenceMixerConfig,
    input_dim: int,
    output_dim: int,
    device: device,
)

Bases: Module

Categorical multi-feature sequence model.

The module consumes pre-encoded inputs shaped (B, T, C_in, K) and emits logits shaped (B, T, C_out, K). Use encode_categorical_values for scalar inputs or the experiment helpers used by training and rollout.

Parameters:

config (SequenceMixerConfig) –

Model architecture.
input_dim (int) –

Number of ordered input features.
output_dim (int) –

Number of ordered targets.
device (device) –

Construction device. Mamba layers require CUDA.

Methods:

forward –

Calculate categorical logits and optional recurrent states.

forward ¶

forward(
    x: Tensor,
    mask: Tensor | None = None,
    *,
    states: dict[str, MambaCarryState] | None = None,
    return_states: bool = False,
) -> Tensor | tuple[Tensor, dict[str, MambaCarryState]]

Calculate categorical logits and optional recurrent states.

Parameters:

x (Tensor) –

Encoded inputs shaped (B, T, C_in, K).
mask (Tensor | None, default: None ) –

Optional valid-row mask shaped (B, T).
states (dict[str, MambaCarryState] | None, default: None ) –

Recurrent states keyed by configured Mamba layer path.
return_states (bool, default: False ) –

Return final state for each stateful layer.

Returns:

Tensor | tuple[Tensor, dict[str, MambaCarryState]] –

Logits shaped (B, T, C_out, K), optionally paired with recurrent
Tensor | tuple[Tensor, dict[str, MambaCarryState]] –

states aligned to the end of the supplied sequence.

Advanced Recurrent State¶

MambaCarryState is an opaque advanced type exposed by stateful model execution. Application code should pass it back only to the same model layer and aligned stream; it should not inspect or share its tensors across unrelated samples.

MambaCarryState `dataclass` ¶

MambaCarryState(
    angle_state: Tensor,
    ssm_state: Tensor,
    k_state: Tensor,
    v_state: Tensor,
)

Opaque recurrent state for one SISO Mamba-3 layer.

State is lane- and window-position-specific; callers must never reuse it across unrelated streams or misaligned window starts.

Categorical Encoding¶

categorical_target_distribution ¶

categorical_target_distribution(
    target: Tensor,
    num_bins: int,
    sigma: float,
    target_ranges: Tensor,
) -> Tensor

Encode scalar targets as distributions over bounded categorical bins.

With sigma=0, probability mass is linearly interpolated between adjacent bins. Positive sigma produces a normalized Gaussian over bin centers. Values outside their target ranges are clamped before encoding.

Parameters:

target (Tensor) –

Values shaped (..., C).
num_bins (int) –

Number of categorical bins.
sigma (float) –

Gaussian width in normalized [0, 1] coordinates.
target_ranges (Tensor) –

Per-channel bounds shaped (C, 2).

Returns:

Tensor –

Float32 distributions shaped (..., C, K).

encode_categorical_values ¶

encode_categorical_values(
    values: Tensor,
    num_bins: int,
    sigma: float,
    value_ranges: Tensor,
) -> Tensor

Encode scalar model inputs as categorical distributions.

Parameters:

values (Tensor) –

Scalar inputs shaped (B, T, C).
num_bins (int) –

Number of categorical bins.
sigma (float) –

Gaussian width in normalized coordinates; zero uses adjacent-bin interpolation.
value_ranges (Tensor) –

Per-channel bounds shaped (C, 2).

Returns:

Tensor –

Encoded values shaped (B, T, C, K) with the input dtype.

Raises:

ValueError –

If input rank or range shape is invalid.

decode_categorical_logits ¶

decode_categorical_logits(
    logits: Tensor, target_ranges: Tensor
) -> Tensor

Decode logits through the expected bin location.

This returns the probability-weighted expectation, not the argmax bin.

Parameters:

logits (Tensor) –

Categorical logits shaped (..., C, K).
target_ranges (Tensor) –

Per-channel output bounds shaped (C, 2).

Returns:

Tensor –

Decoded values shaped (..., C) in output-scaled units.

ML Models¶

SequenceMixerConfig dataclass ¶

LayerConfig dataclass ¶

ResidualConfig dataclass ¶

MambaConfig dataclass ¶

OutputConfig dataclass ¶

build_model ¶

SequenceMixer ¶

forward ¶

Advanced Recurrent State¶

MambaCarryState dataclass ¶

Categorical Encoding¶

categorical_target_distribution ¶

encode_categorical_values ¶

decode_categorical_logits ¶

SequenceMixerConfig `dataclass` ¶

LayerConfig `dataclass` ¶

ResidualConfig `dataclass` ¶

MambaConfig `dataclass` ¶

OutputConfig `dataclass` ¶

MambaCarryState `dataclass` ¶