lgatr.nets.conditional_lgatr.ConditionalLGATr

class lgatr.nets.conditional_lgatr.ConditionalLGATr(num_blocks, in_mv_channels, mv_channels_cond, out_mv_channels, hidden_mv_channels, in_s_channels, s_channels_cond, out_s_channels, hidden_s_channels, attention, crossattention, mlp, primitives=None, dropout_prob=None, norm_elementwise_affine=True, checkpoint_blocks=False, naive_amp=False, compile=False, compile_kwargs=None, activation_memory_budget=None)[source]

Bases: Module

Conditional L-GATr network.

Combines num_blocks ConditionalLGATrBlock modules (geometric self-attention, cross-attention, geometric MLP, residual connections, normalization) with initial and final equivariant linear layers. The condition is expected to be already preprocessed (e.g. by a non-conditional LGATr network).

Parameters:

  • num_blocks (int) – Number of transformer blocks.

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

  • mv_channels_cond (int) – Number of condition 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.

  • s_channels_cond (int) – Number of scalar condition channels. Use 0 for no scalar condition stream.

  • 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.

  • crossattention (CrossAttentionConfig | Mapping) – Cross-attention configuration.

  • mlp (MLPConfig | Mapping) – MLP configuration.

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

  • 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 transformer 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 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, multivectors_cond, scalars=None, scalars_cond=None, attn_kwargs=None, crossattn_kwargs=None)[source]

Forward pass.

Parameters:

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

  • multivectors_cond (Tensor) – Condition multivectors of shape (..., items_cond, mv_channels_cond, 16).

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

  • scalars_cond (Tensor | None) – Optional condition scalars 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 | 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.