---
id: cuequivariance-ops-torch-cu12
version: "0.11.1"
license: Other/Proprietary License (NVIDIA Proprietary Software) # Software License Agreement LICENSE AGREEMENT FOR NVIDIA MATH LIBRARIES SOFTWARE DEVELOPMENT KITS This license agreement(“Agreement”) is a… (full text in the JSON record)
license_treatment: unclear
maintenance: active
---

# cuequivariance-ops-torch-cu12 — cuequivariance-ops-torch - GPU Accelerated Torch Extensions for Equivariant Primitives

License: unclear · Maintenance: active · Downloads: 122.6K/mo

## 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 above — 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

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

- License: Other/Proprietary License (NVIDIA Proprietary Software) # Software License Agreement LICENSE AGREEMENT FOR NVIDIA MATH LIBRARIES SOFTWARE DEVELOPMENT KITS This license agreement(“Agreement”) is a… (full text in the JSON record) (unclear)
- Python support: supports_current
- Install friction: medium
- Maintenance: active
- Downloads: 122.6K/month (top 15,000 on PyPI)
- Known vulnerabilities: none known

## Tags

cuda pytorch kernels equivariant, gpu accelerated neural network primitives, pairformer inference optimization, low-precision tensor operations torch, equivariant deep learning cuda, cuda-kernels, equivariant-networks, inference-optimization

[View on SkillFed](https://skillfed.io/packages/cuequivariance-ops-torch-cu12) · [View on PyPI](https://pypi.org/project/cuequivariance-ops-torch-cu12/)