API Reference

L-GATr Networks

We provide two main L-GATr networks, LGATr as a stack of transformer encoders, and ConditionalLGATr as a stack of transformer decoders. For tasks where conditional inputs are required, you can process the condition with a LGATr and then include this processed condition using a ConditionalLGATr. In addition LGATrSlim and ConditionalLGATrSlim provide more efficient versions of the respective networks using only scalar and vector representations.

lgatr.nets.lgatr.LGATr(num_blocks, ...[, ...])

L-GATr network.

lgatr.nets.conditional_lgatr.ConditionalLGATr(...)

Conditional L-GATr network.

lgatr.nets.slim.LGATrSlim(num_blocks, ...[, ...])

L-GATr-slim network.

lgatr.nets.conditional_slim.ConditionalLGATrSlim(...)

Conditional L-GATr-slim network.

L-GATr Layers

The LGATr and ConditionalLGATr networks have a structure similar to standard transformers. We construct them using variants of the standard transformer layers adapted to the geometric algebra framework.

lgatr.layers.lgatr_block.LGATrBlock(...[, ...])

L-GATr encoder block.

lgatr.layers.conditional_lgatr_block.ConditionalLGATrBlock(...)

L-GATr decoder block.

lgatr.layers.linear.EquiLinear(...[, ...])

Linear layer.

lgatr.layers.attention.self_attention.SelfAttention(...)

L-GATr self-attention.

lgatr.layers.attention.cross_attention.CrossAttention(...)

L-GATr cross-attention.

lgatr.layers.mlp.mlp.GeoMLP(config, primitives)

MLP with a geometric product as the first nonlinear mixing step.

lgatr.layers.mlp.geometric_bilinears.GeometricBilinear(...)

Pin-equivariant bilinear map that constructs new geometric features via geometric products.

lgatr.layers.mlp.nonlinearities.ScalarGatedNonlinearity([...])

Gated nonlinearity on multivectors.

lgatr.layers.layer_norm.EquiLayerNorm([...])

Layer normalization for multivectors (and an optional scalar stream).

lgatr.layers.dropout.GradeDropout([p])

Grade-wise dropout on multivectors (and an optional scalar stream).

L-GATr Primitives

The L-GATr primitives implement the core equivariant operations and are called by the L-GATr layers.

lgatr.primitives.attention

Equivariant attention.

lgatr.primitives.bilinear

Geometric product.

lgatr.primitives.dropout

Grade dropout.

lgatr.primitives.invariants

Invariants: inner product, absolute squared norm, and Pin-invariant utilities.

lgatr.primitives.linear

Linear operations on multivectors, in particular linear basis maps.

lgatr.primitives.normalization

Multivector normalization.

L-GATr Configuration Classes

L-GATr uses dataclass objects to organize less relevant hyperparameters like number of heads or the MLP nonlinearity. The MLPConfig, SelfAttentionConfig, CrossAttentionConfig, and PrimitivesConfig are all arguments for the LGATr/ConditionalLGATr modules.

lgatr.primitives.config.PrimitivesConfig([...])

Symmetry-group and bilinear-layer toggles for an L-GATr model.

lgatr.layers.attention.config.SelfAttentionConfig([...])

Configuration for self-attention.

lgatr.layers.attention.config.CrossAttentionConfig([...])

Configuration for cross-attention.

lgatr.layers.mlp.config.MLPConfig([...])

Geometric-MLP configuration.

Interface to the Geometric Algebra

Before we feed data into L-GATr networks and after we extract results, we have to convert between common scalar/vector objects and multivectors. This is very simple, we still introduce convenience methods for this step. We also include functionality to construct spurions, or reference multivectors, which can be added as extra items or channels to break equivariance at the input level.

lgatr.interface.scalar

Embedding and extracting scalars into multivectors.

lgatr.interface.vector

Embedding and extracting vectors into multivectors.

lgatr.interface.bivector

Embedding and extracting bivectors into multivectors.

lgatr.interface.axialvector

Embedding and extracting axial vectors into multivectors.

lgatr.interface.pseudoscalar

Embedding and extracting pseudoscalars into multivectors.

lgatr.interface.spurions

Tools to include reference multivectors ('spurions') for symmetry breaking.

L-GATr Utilities

Helpers used by the L-GATr networks: a wrapper around torch.compile() for the compile=True constructor path, a warmup_caches() helper that pre-populates the primitive caches for a given (device, dtype), and an autocast decorator that pins inputs to a minimum precision.

lgatr.utils.compile

Helpers for using L-GATr networks with torch.compile().

lgatr.primitives.compile

Helpers for using L-GATr primitives with torch.compile().

lgatr.utils.autocast

Pin inputs to a minimum autocast precision; usable as a decorator.

L-GATr-slim Layers

In addition to the full L-GATr network, we provide a slimmed-down version that uses only scalar and vector representations instead of full multivectors. This approach allows a more efficient implementation while achieving similar performance on all high-energy physics tasks we have tested so far.

lgatr.layers.slim_layers.SlimBlock(...[, ...])

A single block of the L-GATr-slim network.

lgatr.layers.slim_layers.ConditionalSlimBlock(...)

A single block of the conditional L-GATr-slim network.

lgatr.layers.slim_layers.SlimSelfAttention(...)

Self-attention for Lorentz vectors and scalar features.

lgatr.layers.slim_layers.SlimCrossAttention(...)

Cross-attention for Lorentz vectors and scalar features.

lgatr.layers.slim_layers.SlimMLP(v_channels, ...)

Multi-layer perceptron for vector and scalar features.

lgatr.layers.slim_layers.SlimGLU(...[, ...])

Gated linear unit (GLU) for vector and scalar features.

lgatr.layers.slim_layers.SlimLinear(...[, ...])

Linear layer for vector and scalar features.

lgatr.layers.slim_layers.SlimRMSNorm(...[, ...])

Joint RMS normalization over vector and scalar features.

lgatr.layers.slim_layers.SlimDropout(...)

Dropout for vector and scalar features.