cuequivariance-torch

What it is and what it does

cuequivariance-torch is NVIDIA's PyTorch frontend for building geometric neural networks that respect symmetries like rotations and translations. It provides CUDA-accelerated operations on segmented polynomials and tensor products, designed to make equivariant models more data-efficient by encoding physical symmetries directly into the network architecture.

The package wraps optimized CUDA kernels behind a PyTorch-native API, letting you construct networks that automatically preserve equivariance properties. It depends on the core cuequivariance library and is intended for developers building physics-informed or geometry-aware deep learning models where respecting spatial symmetries improves both sample efficiency and model interpretability.

Use it for:

• Build rotation- and translation-invariant neural networks for molecular or crystal structure prediction
• Accelerate training of geometric deep learning models on GPU with native CUDA kernels
• Construct equivariant feature extractors for 3D point cloud or mesh-based computer vision tasks
• Encode physical symmetries into neural network layers for physics simulation or inverse problems

Worth the install?

AI-flagged interpretation of the facts on this page — verify before relying

PyTorch bindings for NVIDIA's CUDA-accelerated library for building equivariant geometric neural networks using segmented polynomials and optimized tensor operations.

Yes, if you are building equivariant geometric neural networks in PyTorch and have access to compatible CUDA hardware. The package is actively maintained, permissively licensed, and backed by NVIDIA. However, it is in Beta state, so evaluate stability for your production timeline. Verify that required CUDA kernel packages are available for your environment before committing.

Install

cuequivariance-torch on PyPI

pip install cuequivariance-torch

uv add cuequivariance-torch

poetry add cuequivariance-torch

Installing cuequivariance-torch

Before you install

Active maintenance with recent release. Low install friction via pure Python wheel. Single runtime dependency on cuequivariance core.

License in practice

Apache 2.0 permissive license allows commercial and derivative use with attribution and license propagation requirements.

Quickstart

pip install cuequivariance-torch
import cuequivariance_torch

Requires Python 3.10 or later and CUDA-capable GPU; CUDA kernel packages (cuequivariance-ops-torch-cu12 or cu13) may be needed for full functionality

Verify before relying

• Whether CUDA kernel packages (cuequivariance-ops-torch-cu12 or cu13) are required for actual use or optional
• Performance characteristics and typical speedup over non-equivariant baselines
• Maturity level beyond Beta status and stability guarantees for production workloads

Package facts

Evidence: cuequivariance_torch-0.11.1-py3-none-any.whl

Tags

More Artificial Intelligence packages

Further reading