Pin-equivariant linear layer. The forward pass maps multivector inputs with shape (..., in_channels, 16) to multivector outputs with shape (.

Geometric MLP configuration. Parameters ---------- mv_channels : iterable of int Number of multivector channels at each layer, fr

Configuration for attention. Parameters ---------- in_mv_channels : int Number of input multivector channels. out_mv_channels

GATr network for a data with a single token dimension. This, together with gatr.nets.axial_gatr.AxialGATr, is the main architecture proposed in o

Grade dropout for multivectors (and regular dropout for auxiliary scalars). Parameters ---------- p : float Dropout probability.

Geometric MLP. This is a core component of GATr's transformer blocks. It is similar to a regular MLP, except that it uses geometric bilinears

Applies rotary position encodings (RoPE) to scalar tensors. References ---------- Jianlin Su et al, "RoFormer: Enhanced Transformer with

Equivariant LayerNorm for multivectors. Rescales input such that `mean_channels |inputs|^2 = 1`, where the norm is the GA norm and the mean g

Equivariant transformer block for GATr. This is the biggest building block of GATr. Inputs are first processed by a block consisting of Laye

GCA-GNN layer as described in D. Ruhe et al. This was described in "Geometric Clifford Algebra Networks" by D. Ruhe et. al, and in private co

Geometric self-attention layer. Constructs queries, keys, and values, computes attention, and projects linearly to outputs. Parameters -

Random Pin transform on a multivector torch.Tensor. Slow, only used for testing purposes. Breaks computational graph.

Axial GATr network for two token dimensions. This, together with gatr.nets.gatr.GATr, is the main architecture proposed in our paper. It com

Baseline transformer block. Inputs are first processed by a block consisting of LayerNorm, multi-head self-attention, and residual connection

GCA-MLP model. As described in D. Ruhe et al, "Geometric Clifford Algebra Networks". Based on layers from the official [cliffordlayers co

Geometric attention layer. This is the main attention mechanism used in GATr. Thanks to the nonlinear features, the scaled-dot-product attent

Geometric bilinear layer. Pin-equivariant map between multivector tensors that constructs new geometric features via geometric products and t

Functionality related to installing the Python code given as a repo.

Gated nonlinearity, where the gate is simply given by the scalar component of the input. Given multivector input x, computes f(x_0) * x, where f

Experiment manager for wall-shear-stress estimation in arteries. The dataset was introduced and is described in Suk et al, "Mesh Convolutional Ne

Baseline axial transformer for data with two token dimensions. Combines num_blocks transformer blocks, each consisting of multi-head self-attenti

Baseline layer norm over all dimensions except the first.

Baseline self-attention layer. Parameters ---------- in_channels : int Number of input channels. out_channels : int N

Baseline transformer. Combines num_blocks transformer blocks, each consisting of multi-head self-attention layers, an MLP, residual connectio

A drop-in replacement for EquiLinear for fast inference. Not trainable. Parameters must be set by hand. Parameters ---------- in_mv_

Geometric cross-attention layer. Constructs queries, keys, and values, computes attention, and projects linearly to outputs. Parameters

Compute (multivector and scalar) queries, keys, and values via multi-query attention. Parameters ---------- config: SelfAttentionConfig

N-body prediction dataset. Loads data generated with generate_nbody_dataset.py from disk. Parameters ---------- filename : str or pa

Experiment manager for n-body prediction. Parameters ---------- cfg : OmegaConf Experiment configuration. See the config folder i

Simulator for the n-body dataset. Each sample consists of the positions of a n particicles (1 star and n - 1 planets), both before and after

Exception to be raise when the training encounters a NaN in loss or model weights.

Artery dataset with added shape id attribute.

Compute (multivector and scalar) queries, keys, and values via multi-head attention. Parameters ---------- config: SelfAttentionConfig

A dataset containing random (S)O(3) elements. Parameters ---------- num : int Size of the dataset (number of rotations) speci

Special installation for the segnn package required. The reason is that rope fails to move all files (utils.py). As we are only interested in

Wrapper that turns a torch Dataset into a torch_geometric-compatible version. Assumes that the wrapped dataset returns tuples (x, y), and wraps t

Wrapper that turns a "plain" GATr wrapper into one that uses xformers.

Wraps around simple baselines (like a Transformer) for the artery experiment. We use the following parameterization of the artery mesh: - the

Wraps around GATr for the arterial wall-shear-stress prediction experiment. Parameters ---------- net : torch.nn.Module GATr mode

Base experiment manager class. To be subclassed by experiment-specific manager classes. Parameters ---------- cfg : OmegaConf

Base GATr wrapper. To be subclassed by experiment-specific wrapper classes. Parameters ---------- net : torch.nn.Module GATr

Toy dataset that maps two points x, y to the vector t = y - x between them.

Wrapper around GATr networks for ConnectPointsDataset.

Warning about a deprecation in the GATr library.

GCA-GNN model as described in D. Ruhe et al. The model was described in "Geometric Clifford Algebra Networks" by D.Ruhe et al., and in privat

A simple baseline MLP. Flattens all dimensions except batch and uses GELU nonlinearities.

Compute queries, keys, and values via multi-head attention. Parameters ---------- in_channels : int Number of input channels.

Compute queries, keys, and values via multi-query attention. Parameters ---------- in_channels : int Number of input channels.

Wraps around simple baselines (MLP or Transformer) for the n-body prediction experiment. Parameters ---------- net : torch.nn.Module

Wraps around GATr for the n-body prediction experiment. Parameters ---------- net : torch.nn.Module GATr model that accepts input

Wraps around GCA-MLP and GCA-GNN baselines for the n-body experiment. Parameters ---------- net : torch.nn.Module GCAN model that

Wraps around the SE3-Transformer baseline for the n-body prediction experiment. Parameters ---------- net : torch.nn.Module SE3-T

Wraps around the SEGNN baseline for the n-body prediction experiment. Parameters ---------- net : torch.nn.Module SEGNN model tha

Toy dataset that maps two points X, Y to the Euclidean distance |X - Y|.

Wrapper around GATr networks for PointsDistanceDataset.

SEGNN baseline. We restrict ourselves to vector and scalar inputs, as those were sufficient for our experiments. References ----------

Toy dataset that maps a point x and a translation vector t to the translated point t+x.

Wrapper around GATr networks for TranslatePointDataset.

Toy dataset that maps a translation vector t to its Euclidean norm |t|.

Wrapper around GATr networks for TranslationNormDataset.

Serves as a `torch.compile`-compatible replacement for `@functools.cache()`.