lgatr.nets.lgatr.LGATr

class lgatr.nets.lgatr.LGATr(num_blocks, in_mv_channels, out_mv_channels, hidden_mv_channels, in_s_channels, out_s_channels, hidden_s_channels, attention, mlp, primitives=None, reinsert_mv_channels=None, reinsert_s_channels=None, dropout_prob=None, norm_elementwise_affine=True, checkpoint_blocks=False, naive_amp=False, compile=False, compile_kwargs=None, activation_memory_budget=None)

Bases: Module

L-GATr network.

Combines num_blocks LGATrBlock modules (geometric self-attention, geometric MLP, residual connections, normalization) with initial and final equivariant linear layers.

Inputs have shape (..., items, in_mv_channels, 16); outputs have shape (..., items, out_mv_channels, 16); hidden representations have shape (..., items, hidden_mv_channels, 16) (and similar for the optional scalar stream).

Parameters:

• num_blocks (int) – Number of transformer blocks.

• in_mv_channels (int) – Number of input multivector channels.

• out_mv_channels (int) – Number of output multivector channels.

• hidden_mv_channels (int) – Number of hidden multivector channels.

• in_s_channels (int) – Number of scalar input channels. Use 0 for no scalar inputs.

• out_s_channels (int) – Number of scalar output channels. Use 0 for no scalar outputs.

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

• attention (SelfAttentionConfig | Mapping) – Self-attention configuration (see SelfAttentionConfig).

• mlp (MLPConfig | Mapping) – MLP configuration (see MLPConfig).

• primitives (PrimitivesConfig | Mapping | None) – LGATr primitives configuration. Accepts a PrimitivesConfig instance, a dict, or None (uses defaults).

• reinsert_mv_channels (tuple[int, ...] | None) – If not None, specifies multivector channels that will be reinserted in every attention layer.

• reinsert_s_channels (tuple[int, ...] | None) – If not None, specifies scalar channels that will be reinserted in every attention layer.

• dropout_prob (float | None) – Dropout probability.

• norm_elementwise_affine (bool) – Whether the block EquiLayerNorm instances learn an affine gain.

• checkpoint_blocks (bool) – Whether to use gradient checkpointing for the blocks. Saves memory at the cost of speed.

• 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 multivector stream and metric contractions stay fp32 while the scalar GEMMs run in bf16.

• compile (bool) – Whether to wrap the model with torch.compile(). Primitive caches are warmed automatically whenever the model is moved or cast (.to(), .cuda(), .float(), etc.), so the captured graph is free of host-to-device copies.

• 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(multivectors, scalars=None, **attn_kwargs)

Forward pass.

Parameters:

• multivectors (Tensor) – Input multivectors of shape (..., items, in_mv_channels, 16).

• scalars (Tensor | None) – Optional input scalars of shape (..., items, in_s_channels).

• **attn_kwargs – Optional keyword arguments forwarded to attention.

Return type:

tuple[Tensor, Tensor | None]

Returns:

• outputs_mv – Output multivectors of shape (..., items, out_mv_channels, 16).

• outputs_s – Output scalars of shape (..., items, out_s_channels), or None if out_s_channels == 0.