cuequivariance-ops-torch (CUDA 12)

0.10.0 verified Mon Apr 27 auth: no python

GPU-accelerated PyTorch extensions for equivariant primitives using CUDA 12. Part of the cuEquivariance ecosystem, providing high-performance implementations for e3nn and similar libraries. Version 0.10.0, actively developed.

pip install cuequivariance-ops-torch-cu12

Common errors

error ModuleNotFoundError: No module named 'cuequivariance_ops_torch' ↓

cause Typo or wrong package import path. The correct import is from 'cuequivariance_ops_torch', not 'cuequivariance_ops'.

fix

Change import to: from cuequivariance_ops_torch import cuEquivarOps

error RuntimeError: CUDA error: no kernel image is available for execution on the device ↓

cause CUDA version mismatch between PyTorch, cuEquivariance ops, and GPU driver.

fix

Ensure CUDA 12 compatible versions of all libraries: pip install torch==2.1.0+cu121 cuequivariance-ops-torch-cu12

error ImportError: cannot import name 'set_irreps' from 'cuequivariance_ops_torch' ↓

cause API changed in version 0.10.0; 'set_irreps' function replaced by 'cuEquivarOps' class constructor.

fix

Use: ops = cuEquivarOps(irreps) instead of ops = set_irreps(irreps)

Warnings

breaking Version 0.10.0 may have breaking changes from 0.9.x. Check migration guide for changes in API signatures (e.g., renaming of 'set_irreps' to 'cuEquivarOps' constructor argument). ↓

fix Review changelog; update imports and method calls accordingly.

gotcha The package is specific to CUDA 12. Install the correct variant: cuequivariance-ops-torch-cu12 (CUDA 12), cuequivariance-ops-torch-cu11 (CUDA 11), or CPU-only version (cuEquivariance-ops-torch). Mixing versions leads to runtime errors. ↓

fix Match package name to your CUDA version: pip install cuequivariance-ops-torch-cu12 for CUDA 12.

gotcha Requires PyTorch compiled with CUDA. If PyTorch was installed from CPU-only wheel, cuEquivariance ops will fail at runtime. ↓

fix Install PyTorch with CUDA support: pip install torch --index-url https://download.pytorch.org/whl/cu121 (adjust CUDA version).

Imports

cuEquivarOps

wrong

from cuequivariance_ops import cuEquivarOps

correct

from cuequivariance_ops_torch import cuEquivarOps

Package name includes '_torch' suffix

set_irreps

from cuequivariance_ops_torch import set_irreps

Direct import from ops module

Quickstart

Initialize cuEquivarOps for given irreps and apply forward pass.

import torch
from cuequivariance_ops_torch import cuEquivarOps

# Define a simple irreducible representation
irreps = '0e + 1o'
ops = cuEquivarOps(irreps, math_dtype=torch.float32)

# Create random input signal
x = torch.randn(4, 3)  # batch=4, vector size 3
output = ops.forward(x)
print(output.shape)