NVIDIA cuEquivariance Documentation#

cuEquivariance is an NVIDIA Python library designed to facilitate the construction of high-performance equivariant neural networks using segmented tensor products. cuEquivariance provides a comprehensive API for describing segmented polynomials made out of segmented tensor products and optimized CUDA kernels for their execution. Additionally, cuEquivariance offers bindings for both PyTorch and JAX, ensuring broad compatibility and ease of integration.

Equivariance is the mathematical formalization of the concept of “respecting symmetries.” Robust physical models exhibit equivariance with respect to rotations and translations in three-dimensional space. Artificial intelligence models that incorporate equivariance are often more data-efficient.

An introduction to group representations can be found in the page Groups and Representations.

Open Source#

cuEquivariance frontend is open-source and available on GitHub under the Apache 2.0 license.

Installation#

The easiest way to install cuEquivariance is from PyPi using pip.

# Choose the frontend you want to use
pip install cuequivariance-jax
pip install cuequivariance-torch
pip install cuequivariance  # Installs only the core non-ML components

# CUDA kernels
pip install cuequivariance-ops-torch-cu11
pip install cuequivariance-ops-torch-cu12
pip install cuequivariance-ops-jax-cu12

Requirements#

cuequivariance-ops-torch-* packages are available for Linux x86_64/aarch64 and require PyTorch 2.4.0 or later. aarch64 is only available for Python 3.12.

cuequivariance-ops-jax-cu12 package is available for Linux x86_64/aarch64 and requires JAX 0.5.0 or later.

Organization#

cuEquivariance is split into three packages:

import cuequivariance as cue
# All the non-ML components

import cuequivariance_jax as cuex
# For the JAX implementations

import cuequivariance_torch as cuet
# For the PyTorch implementations

Most tensor products are defined using a hierarchy of components:

cue.EquivariantPolynomial - Encapsulates the cue.Rep for each input and the output
cue.SegmentedPolynomial - Included in the EquivariantPolynomial, composed of one or several tensor products
cue.SegmentedTensorProduct - Defines the specific tensor product operations

This descriptor can then be used in two ways:

In PyTorch: Create a cuet.SegmentedPolynomial module for use in models
In JAX: Execute using cuex.equivariant_polynomial or cuex.segmented_polynomial

Tutorials#

Groups and Representations
Data Layouts
Segmented Polynomials
Segmented Tensor Product

API Reference#

What’s New#

Release Notes