Compute basis features for queries and keys in the geometric SDP attention. Parameters ---------- device: torch.device Device.

Builds graph for a single sample.

Builds graph for a full batch.

Computes accelerations for a set of point masses according to Newtonian gravity.

Constructs a tensor for the dual operation. Parameters ---------- device : torch.device Device dtype : torch.dtype Dt

Computes position embeddings and stores them. The position embedding is computed along dimension `item_dim` of tensor `inputs` and is

Constructs a matrix that computes multivector grade involution. Parameters ---------- device : torch.device Device dtype : to

Computes prefactors for the initialization. See self.reset_parameters().

Constructs everything we need to compute norm(equi_norm(x,y)) in a memory-efficient way. Parameters ---------- threshold : float

Constructs input features that will be reinserted in every attention layer.

Creates the model from the config. Returns ------- model : torch.nn.Module A randomly initialized model, followin

Creates optimizer and scheduler. Returns ------- optim : torch.optim.Optimizer Adam optimizer for the parameters

Computes einsum with a given contraction path.

Creates source and destination indices for a fully connected graph.

Provides shape-based caching of the optimal contraction path.

Initializes device, dtype, and attention implementation.

Weight initialization for maps to multivector outputs.

Initializes matplotlib's rcparams to look good.

Weight initialization for maps to multivector outputs.

Initialize experiment folder (and plenty of other initialization thingies).

Creates experiment folder.

Initializes all things related to MLflow tracking.

Loads data from file and converts to input and output tensors.

Log to MLflow. Parameters ---------- loss : torch.Tensor Loss metrics : dict with str keys and float valu

PGA-MLP code as sent by David Ruhe.

MLP block.

Use rope to move files.

Optimizer step and gradient norm clipping.

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

Utility function that "rotates" a tensor, as required for rotary embeddings.

Samples planet velocities around those that give stable circular orbits.

Stores the config in the experiment folder and tracks it with mlflow.

A helper function to isolate the graph-breaking parts of the attention (cf. decorator). TODO: This function can be dissolved once we get expand_p

Performs random E(3) transformations and permutations on given positions / velocities.

Evolves an initial state under Newtonian equations of motions.

Sorts files for priority refactoring. This is necessary as else rope will get this wrong in some cases: It can, for example, make a d

Everything that may happen per step.

Adds an `__init__.py` file for ever subfolder of checkout path (if missing).

Adds a generic `pyproject.toml` file to the project root.

Applies patch from a patch file, if provided upon instantiation.

Block diagonally stack tensors along dimensions dim1 and dim2.

Construct attention mask. Parameters ---------- inputs : torch.Tensor Raw inputs to wrapped network. mv :

Checks whether a callable computes the reversal of a multivector.

Checks whether a callable computes the geometric product.

Checks whether a callable computes the grade involution of a multivectors.

Checks whether a callable computes the outer product.

Checks whether a callable computes the reversal of a multivector.

Checks out the commit hash provided upon instantiation.

Clones the git repository from the clone URL provided upon instantiation.

Constructs a bool array for the inner product calculation. The inner product of MVs is <~x y>_0, i.e. take the grade-0 component of the geometric

Computes the nominal FLOPs per token for an EquiLinear layer. We assume: - the number of tokens are large, so the token-independent contract

Embeds the pose of a 3D object PGA multivectors. The pose is given by a position and a orientation quaternion. The latter is assumed to follow

Embeds the pose of a 3D object in two-vec parameterization into the PGA. The inputs consist of the position of an object and two vectors that spe

Embeds raw inputs into the geometric algebra (+ scalar) representation. To be implemented by subclasses. Parameters --------

Embeds raw inputs into the geometric algebra (+ scalar) representation. To be implemented by subclasses. Parameters --------

Evaluates self.model on all eval datasets and logs the results.

Extracts 3D position and quaternion orientation from the multivector embedding of an object. This function inverts embed_3d_object(). The ex

Extracts 3D position and two-vec orientation from the multivector embedding of an object. This inverts embed_3d_object_two_vec(). Parameters

Embeds raw inputs into the geometric algebra (+ scalar) representation. To be implemented by subclasses. Parameters --------

Extracts raw outputs from the GATr multivector + scalar outputs. To be implemented by subclasses. Parameters ----------

Extracts predicted positions from PGA multivectors.

Extract ray in Plücker coordinates from a multivector. Plücker coords are (v, o x v) for ray through o in direction v. Args: multive

Extracts the normal on an reflection plane from a multivector. References ---------- Leo Dorst, "A Guided Tour to the Plane-Based Geometr

DEPRECATED: Given a multivector, extract 3D translation vector from the bivector components. Note bene: this function is NOT equivariant, unless

Flattens a nested dictionary with str keys.

Pin-equivariant gated GeLU nonlinearity without division. Given multivector input x and scalar input gates (with matching batch dimensions), comp

Pin-equivariant gated GeLU nonlinearity with division. Given multivector input x and scalar input gates (with matching batch dimensions), compute

Computes einsum with a given contraction path (which is ignored when using compilation). Cf. `enable_cached_einsum` for more context.

Generates a canonical set of datasets for the n-body problem, stores them in `path`.

Given either a tensor or a list of tensors or a dict of tensors, returns the batchsize.

Gets list of all Python packages to install.

Gets parity of a clifford multivector. Given a clifford multivector, returns True if it is pure-odd-grade, False if it is pure-even grade, an

Extracts version information. Inspired by item "1." @ https://packaging.python.org/guides/single-sourcing-package-version/.

Performs all steps necessary up to and including installation of the package.

Refactors flat Python projects into installable form.

EquiLinear fixture for testing

Loads a model checkpoint from disk. Parameters ---------- checkpoint : None or str or pathlib.Path Path to checkp

Entry point for artery experiment.

Entry point for n-body experiment.

Entry point for n-body dataset generation.

Constructs an edge index for fully connected graph.

Creates a PyG-style edge index for a fully connected graph of `num_nodes` nodes.

Return dtype with maximum precision. Compatible with compilation.

Moves Python files and folders into folder with the project's name, adapts imports.

Moves package code to `src` folder. That is, we move the code from `<checkout_base_dir> / project-name / project-name` to `<checkout

Installs the refactored package using `pip`.

Removes all packaging files.

Initializes the weights of the layer. Parameters ---------- initialization : {"default", "small", "unit_scalar", "almost_unit

Moves all PGA instances in GCAN component to device / dtype.

Samples from gravity problem. Parameters ---------- num_samples : int Number of samples. num_planets : in

Samples from the Pin(3,0,1) group as a product of reflections.