NVIDIA cuEquivariance Documentation#
cuEquivariance is a 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 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 for different CUDA versions
pip install cuequivariance-ops-torch-cu11
pip install cuequivariance-ops-torch-cu12
Supported platforms#
cuequivariance-ops-torch-* packages are only available for Linux x86_64.
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 the cue.EquivariantTensorProduct class, which encapsulates the cue.Irreps and cue.IrrepsLayout for each input and the output. It also includes one or more instances of cue.SegmentedTensorProduct, which define the tensor product operations.
This descriptor is then used to create a cuet.EquivariantTensorProduct module, which can be used in PyTorch models. Or used to execute the tensor product operations using cuex.equivariant_tensor_product in JAX.