lgatr.nets.conditional_slim.ConditionalLGATrSlim

class lgatr.nets.conditional_slim.ConditionalLGATrSlim(num_blocks, in_v_channels, v_channels_cond, out_v_channels, hidden_v_channels, in_s_channels, s_channels_cond, out_s_channels, hidden_s_channels, num_heads, nonlinearity='gelu', nonlinearity_v='sigmoid', mlp_ratio=2, attn_ratio=1, num_layers_mlp=2, dropout_prob=None, norm_elementwise_affine=True, checkpoint_blocks=False, naive_amp=False, compile=False, compile_kwargs=None, activation_memory_budget=None)

Bases: Module

Conditional L-GATr-slim network.

Stacks num_blocks ConditionalSlimBlock modules between initial and final SlimLinear layers.

Parameters:

• num_blocks (int) – Number of Lorentz-transformer blocks.

• in_v_channels (int) – Number of input vector channels.

• v_channels_cond (int) – Number of conditional vector channels.

• out_v_channels (int) – Number of output vector channels.

• hidden_v_channels (int) – Number of hidden vector channels.

• in_s_channels (int) – Number of input scalar channels.

• s_channels_cond (int) – Number of conditional scalar channels.

• out_s_channels (int) – Number of output scalar channels.

• hidden_s_channels (int) – Number of hidden scalar channels.

• num_heads (int) – Number of attention heads.

• nonlinearity (str) – Nonlinearity for the MLP layers.

• nonlinearity_v (str | None) – Optional override for the vector-path gate nonlinearity in every GLU. None falls back to nonlinearity.

• mlp_ratio (int) – Expansion ratio for MLP hidden channels.

• attn_ratio (int) – Expansion ratio for attention hidden channels.

• num_layers_mlp (int) – Number of layers in each MLP (must be >= 2).

• dropout_prob (float | None) – Dropout probability.

• norm_elementwise_affine (bool) – Whether the SlimRMSNorm instances learn per-channel gains.

• checkpoint_blocks (bool) – Whether to use gradient checkpointing for the blocks.

• naive_amp (bool) – Whether to bypass the fp32 precision islands so the whole forward runs in the surrounding autocast dtype (e.g. bf16). When False (default), under autocast the vector stream and metric contractions stay fp32 while the scalar GEMMs run in bf16.

• compile (bool) – Whether to wrap the model with torch.compile().

• compile_kwargs (Mapping | None) – Dict forwarded verbatim to torch.compile() (via lgatr.utils.compile.compile_model()) when compile=True (e.g. mode, dynamic, fullgraph). Omitted keys fall back to torch’s own defaults.

• activation_memory_budget (float | None) – Fraction in [0, 1] forwarded to lgatr.utils.compile.compile_model() when compile=True. None (the default) leaves torch’s global setting untouched. Setting 1.0 recomputes only cheap pointwise/reduction ops in the backward pass (torch default); lower values let the partitioner also recompute compute-intensive ops, ranked by memory-saved-per-FLOP, trading backward FLOPs for a smaller activation-memory peak. Smaller values (down to ~0.3) reduce the activation-memory peak at a modest backward-compute cost.

forward(vectors, vectors_cond, scalars, scalars_cond, attn_kwargs=None, crossattn_kwargs=None)

Forward pass.

Parameters:

• vectors (Tensor) – Lorentz vectors of shape (..., items, in_v_channels, 4).

• vectors_cond (Tensor) – Condition Lorentz vectors of shape (..., items_cond, v_channels_cond, 4).

• scalars (Tensor) – Scalar features of shape (..., items, in_s_channels).

• scalars_cond (Tensor) – Condition scalar features of shape (..., items_cond, s_channels_cond).

• attn_kwargs (dict | None) – Optional keyword arguments forwarded to self-attention.

• crossattn_kwargs (dict | None) – Optional keyword arguments forwarded to cross-attention.

Return type:

tuple[Tensor, Tensor]

Returns:

• outputs_v – Lorentz vectors of shape (..., items, out_v_channels, 4).

• outputs_s – Scalar features of shape (..., items, out_s_channels).