cuequivariance-ops-torch-cu12

What it is and what it does

cuequivariance-ops-torch-cu12 is a PyTorch extension that wraps optimized CUDA kernels for equivariant neural network operations. It is designed primarily as a lower-level component accessed through cuEquivariance, though the kernels can be used directly. The package exposes operations as both torch.nn.Module objects (for training and model export) and torch.library custom operators (for torch.compile tracing and TensorRT inference).

The package includes specialized primitives for Pairformer inference, such as attention pair bias masking and combined triangle multiplication. These operations are tuned for specific configurations (BF16 precision, D=256, H=16, sequence lengths 384 or 512) and use native low-precision tensor-core products with FP32 accumulation. It requires CUDA 12, torch, and scipy, and is distributed as platform-specific wheels for x86_64 and aarch64 Linux only.

Use it for:

• Accelerate Pairformer inference in equivariant models using optimized CUDA kernels
• Export equivariant neural network models via torch.export for deployment with TensorRT on NVIDIA GPUs
• Integrate low-precision (BF16/FP16) equivariant operations into production inference pipelines
• Access raw CUDA kernels for custom equivariant layer implementations beyond the standard API

Worth the install?

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

Provides CUDA-accelerated PyTorch kernels and operators for equivariant neural network primitives, including specialized Pairformer inference operations optimized for low-precision tensor computations.

Yes, with conditions. Install if you are building equivariant neural networks on NVIDIA GPUs and need production-grade inference performance, or if you are already using cuEquivariance and want direct kernel access. Do not install if you lack CUDA 12 hardware, need cross-platform support, or require an open-source-compatible license. The proprietary restrictions and hardware lock-in are significant trade-offs for the performance gains.

Install

cuequivariance-ops-torch-cu12 on PyPI

pip install cuequivariance-ops-torch-cu12

uv add cuequivariance-ops-torch-cu12

poetry add cuequivariance-ops-torch-cu12

Installing cuequivariance-ops-torch-cu12

Before you install

Medium install friction due to platform-specific wheels (x86_64 and aarch64 Linux only, CUDA 12 required). Active maintenance with release 7 days ago. Requires torch and scipy as runtime dependencies.

License in practice

NVIDIA proprietary license with significant restrictions: non-transferable, no sublicensing, reverse engineering prohibited, and cannot be used in open-source contexts. Distribution requires material additional functionality beyond the SDK. Use in critical applications is explicitly prohibited and indemnified.

Quickstart

pip install cuequivariance-ops-torch-cu12

import cuequivariance_ops_torch
import torch

# Access kernels as torch.nn.Module or torch.library operators
# See module docstrings for complete input contracts

Requires CUDA 12 and NVIDIA GPU; platform-specific wheels for x86_64 or aarch64 Linux only

Verify before relying

• Performance characteristics and speedup claims for specific model architectures or batch sizes
• Compatibility with specific PyTorch versions beyond Python 3.10–3.14 support
• Whether the package works on systems without NVIDIA GPUs or only on NVIDIA hardware

Package facts

Evidence: cuequivariance_ops_torch_cu12-0.11.1-cp310-cp310-manylinux_2_26_aarch64.manylinux_2_28_aarch64.whl; cuequivariance_ops_torch_cu12-0.11.1-cp310-cp310-manylinux_2_27_x86_64.manylinux_2_28_x86_64.whl; cuequivariance_ops_torch_cu12-0.11.1-cp311-cp311-manylinux_2_26_aarch64.manylinux_2_28_aarch64.whl; cuequivariance_ops_torch_cu12-0.11.1-cp311-cp311-manylinux_2_27_x86_64.manylinux_2_28_x86_64.whl; cuequivariance_ops_torch_cu12-0.11.1-cp312-cp312-manylinux_2_26_aarch64.manylinux_2_28_aarch64.whl; cuequivariance_ops_torch_cu12-0.11.1-cp312-cp312-manylinux_2_27_x86_64.manylinux_2_28_x86_64.whl; cuequivariance_ops_torch_cu12-0.11.1-cp313-cp313-manylinux_2_26_aarch64.manylinux_2_28_aarch64.whl; cuequivariance_ops_torch_cu12-0.11.1-cp313-cp313-manylinux_2_27_x86_64.manylinux_2_28_x86_64.whl; cuequivariance_ops_torch_cu12-0.11.1-cp314-cp314-manylinux_2_26_aarch64.manylinux_2_28_aarch64.whl; cuequivariance_ops_torch_cu12-0.11.1-cp314-cp314-manylinux_2_27_x86_64.manylinux_2_28_x86_64.whl

Tags

More Artificial Intelligence packages

Further reading