Send to device. Overwritten to also move the CliffordAlgebra object, which is not an nn.Module.

Computes the geometric product f(x,y) = xy. Parameters ---------- x : torch.Tensor with shape (..., 16) First input multivector.

Checks whether a callable is equivariant with respect to the Pin(3,0,1) or Spin(3,0,1) group. The callable can have an arbitray number of multive

Embeds 3D points in multivectors. We follow the convention used in the reference below and map points to tri-vectors. References -------

Make tensor n-dimensional, group extra dimensions in first.

Transforms a list of multivector objects to a torch.Tensor.

Given a multivector, extract any potential 3D point from the trivector components. Nota bene: if the output is interpreted a regular R^3 point,

Embeds a 3D translation in multivectors. In our convention, a translation vector is embedded into a combination of the scalar and bivector co

Embeds a scalar tensor into multivectors. Parameters ---------- scalars: torch.Tensor with shape (..., 1) Scalar inputs. Ret

Computes the reversal of a multivector. The reversal has the same scalar, vector, and pseudoscalar components, but flips sign in the bivector

Embeds an (oriented plane) in the PGA. Following L. Dorst, the plane is represent as PGA vectors. References ---------- Leo Dorst, "

Extracts the normal on an oriented plane from a multivector. Currently, this function does *not* extract a support point for the plane (or the di

Utility function that samples a list of multivectors.

Logs metric or param to MLflow. Catches exceptions due to locked or unavailable DB and then retries. Parameters ---------- key : str

Computes the dual of `inputs` (non-equivariant!). See Table 4 in the reference. References ---------- Leo Dorst, "A Guided Tour to t

Embeds 3D rotations in multivectors. We follow the convention used in Leo Dorst, "A Guided Tour to the Plane-Based Geometric Algebra PGA", an

Checks whether an action should be performed at a given step and frequency. If `include_fractional` is given, the check also returns True when

High-level training function.

Embeds the reflection on a plane in multivectors. Following L. Dorst, this is represented as a PGA vector. References ---------- Leo

Given a multivector x, returns |x_{123}| - 1. Put differently, this is the deviation of the norm of a pseudoscalar component from 1. This can

Extracts scalar components from multivectors. Parameters ---------- multivectors: torch.Tensor with shape (..., 16) Multivector i

Computes torch.einsum with contraction path caching if enabled (and compilation is not used). Cf. `enable_cached_einsum` for more context.

Samples from n-body simulator and stores the results at `filename`.

Multivector dropout, dropping out grades independently. Parameters ---------- x : torch.Tensor with shape (..., 16) Input data.

Computes the inner product of multivectors f(x,y) = <x, y> = <~x y>_0. Sums over the 16 multivector dimensions. Equal to `geometric_product(

Computes the GA norm of an input multivector. Equal to sqrt(inner_product(x, x)). NOTE: this primitive is not used widely in our architectur

Given a dataloader, computes all relevant metrics. Can be adapted by subclasses. Parameters ---------- dataloader : torch.uti

Constructs basis elements for Pin(3,0,1)-equivariant linear maps between multivectors. This function is cached. Parameters ----------

Creates an Irreps object for SEGNNModel given a number of O(3) vectors and scalars.

GATrBlock fixture for testing

Casts an object as MLPConfig.

Selects whether to use caching of optimal paths in einsum contraction computations. When using torch.compile (torch==2.2.1), if we specify the pr

Equivariant geometric attention based on scaled dot products and nonlinear aux features. This is the main attention mechanism used in GATr. Thank

Extracts the first output of a tuple of multiple outputs. If only one output is present, returns that. This is convenient for equivariance c

Computes the grade involution of a multivector. The reversal has the same scalar, bivector, and pseudoscalar components, but flips sign in the

Computes the outer product `f(x,y) = x ^ y`. Parameters ---------- x : torch.Tensor with shape (..., 16) First input multivector.

Returns a unit quaternion, uniformly sampled from the 3-sphere.

Executes command in a subprocess, and in case of errors, prints stdout/stderr and fails.

Loads basis elements for Pin-equivariant bilinear maps between multivectors. Parameters ---------- kind : {"gp", "outer"} Filenam

Loads dataset. To be implemented by subclasses. Parameters ---------- tag : str Dataset tag, like "train", "val",

Creates a data loader. Parameters ---------- dataset : torch.nn.utils.data.Dataset Dataset. batch_size :

Checks whether a callable is invariant with respect to the Pin(3,0,1) or Spin(3,0,1) group. Parameters ---------- function: Callable

Constructs a reference vector for the equivariant join. Parameters ---------- reference : Tensor with shape (..., 16) or {"data", "canoni

Embeds a pseudoscalar tensor into multivectors. Parameters ---------- pseudoscalars: torch.Tensor with shape (..., 1) Pseudoscala

Computes the equivariant join. ``` equivariant_join(x, y; reference) = reference_123 * dual( dual(x) ^ dual(y) ) ``` This function u

Extracts pseudoscalar components from multivectors. Nota bene: when the output is interpreted as a scalar, this function is only equivariant

Execute (vanilla) scaled dot-product attention. Dynamically dispatch to xFormers if attn_mask is an instance of xformers.ops.AttentionBias or

Build 5D vector whose inner product with another such vector computes the squared distance. Parameters ---------- tri: Tensor Bat

Computes einsum whilst caching the optimal contraction path. Inspired by upstream https://github.com/pytorch/pytorch/blob/v1.13.0/torch/funct

Casts a single input to at least 32-bit precision.

Casts a single output to desired precision.

Constructs a kernel for the join operation. The kernel is such that join(x, y)_i = einsum(kernel_ijk, x_j, x_k). For now, we do this in the

Constructs a matrix that computes multivector reversal. Parameters ---------- device : torch.device Device dtype : torch.dtyp

Model forward pass. To be implemented by subclasses. Parameters ---------- data : tuple of torch.Tensor Data batc

Initializes logging.

Data preparation during training loop, e.g. to move data to correct device and dtype.

Silences other loggers.

Assert all values in sequence are equal.

Construct attention mask from pytorch geometric batch. Parameters ---------- inputs : torch_geometric.data.Batch Data batch.

Transform EquiLinear module to equivalent CompiledLinear module. Used for fast inference, not trainable! Parameters ---------- equi_

Replace all EquiLinear submodules in module with compiled version. Parameters ---------- model : nn.Module Model whose EquiLinear

Create self.model according to the specification in self.cfg.

Computes the equivariant join, using the efficient implementation. ``` equivariant_join(x, y; reference) = reference_123 * dual( dual(x) ^ du

Represent the n-body initial state in PGA multivectors. Masses are represented as scalars, positions as trivectors, and velocities as bivectors

Embed ray in Plücker coordinates as a multivector. Plücker coords are (v, o x v) for ray through o in direction v. Args: pluecker_ra

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

Pin-equivariant linear map f(x) = sum_{a,j} coeffs_a W^a_ij x_j. The W^a are 9 pre-defined basis elements. Parameters ---------- x :

Expand tensors to largest, optionally excluding some axes.

Computes the equivariant join, using the explicit, but slow, implementation. ``` equivariant_join(x, y; reference) = reference_123 * dual( du

Extracts the represented rotation quaternion from the multivector's bivector components. Parameters ---------- multivector : torch.Tensor

Compute cross attention. Parameters ---------- multivectors_kv : torch.Tensor with shape (..., num_items_kv, channels_in, 16)

Gets CUDA if available, CPU else.

Projects an input tensor to the individual grades. The return value is a single tensor with a new grade dimension. NOTE: this primitive is n

Computes the norm of the join, `|join(x,y)|`, in a single operation. Optionally: - computes the squared norm instead of the norm (when `squar

Return dtype with maximum precision. Cached. Compatible with compilation.

Checks whether models have close weights.

Shorthand to transform a numpy array to a PGA multivector.

Equivariant geometric attention based on scaled dot products and the equivariant join. Expects both multivector and scalar queries, keys, and val

Postprocesses the outputs extracted from the GA representation. To be implemented by subclasses, optionally (by default, no postprocessing is

Normalises q and maps to group matrix. Assumes [x y z w] quaternion format with ij=k (Hamilton) convention.

Samples log-uniformly from (min_, max_).

Save model in experiment folder. Parameters ---------- filename : None or str Filename to save the model to.

Seeds all known sources of pseudo-randomness in our stack.

Transforms a torch.Tensor to a list of multivector objects.

Updates node states with aggregated messages.

Runs validation loop, logs results, and may store state dict for early stopping. Parameters ---------- dataloader : torch.uti

Scaled dot-product attention.

Attention block.

SEGNN feature engineering for n-body experiments. Constructs node features (position relative to mean position, velocity embedding, absolute