/bionemo-cuequivariance-torch

name: cuequivariance-torch
description: Use when building or debugging cuEquivariance PyTorch tensor products or equivariant GNNs, especially CPU references, layouts, indexing, and gradients.
license: Apache-2.0
metadata:
  author: Mario Geiger <mgeiger@nvidia.com>
  tags: [pytorch, equivariance, tensor-products]

cuEquivariance in PyTorch

Purpose

Execute `cuequivariance` (`cue`) descriptors with `cuequivariance_torch` (`cuet`): CPU references, CUDA backends, equivariant layers, and autograd. Preserve the project's environment.

Inputs

Required inputs: operation, SO3/O3 irreps, layout, device, dtype, and weight ownership; infer from the request or code. Graphs also need source/destination counts and indices. Optional inputs: backend, performance, export targets.

Instructions

1. Probe once and reuse working PyTorch. **Never recursively delete directories during setup or recovery**, including partial temporary environments; use a fresh unique path when isolation is needed. 2. Set `method="naive"` at construction for CPU. Match parameters, features, weights, and references in device and dtype. Use float64 for small numerical and gradient checks. 3. Specify layout and weight ownership explicitly. Derive sizes and output order from the module/descriptor; do not guess multiplicities from a filter. 4. Finish installation before running the saved script with the completion recipe below. A session/task ID, empty output, or a file write is not completion. Poll that session until it exits and inspect its results. 5. For an API mismatch, inspect the installed signature/source once. Adapt or explain the requirement; preserve the requested library.

Prerequisites and environment reuse

Probe imports independently with the project's Python; a missing frontend must not hide working PyTorch:

import importlib
import sys

print("python:", sys.executable, sys.version.split()[0])
available, import_errors = {}, {}
for name in ("torch", "cuequivariance", "cuequivariance_torch"):
    try:
        module = importlib.import_module(name)
        available[name] = module
        print(name, getattr(module, "__version__", "unknown"), module.__file__)
    except (ImportError, OSError) as error:
        import_errors[name] = f"{type(error).__name__}: {error}"
        print(name, import_errors[name])
if "torch" in available:
    torch = available["torch"]
    print("CUDA build:", torch.version.cuda, "available:", torch.cuda.is_available())

Choose the next step from the probe:

• Imports work: retain the interpreter; install nothing.
• Only cuEquivariance is missing: install `cuequivariance-torch` in that interpreter; it does not require replacing PyTorch.
• PyTorch is missing: check the project interpreter (e.g. `.venv/bin/python`). If the requested installation is inaccessible, say so; a new environment cannot verify it.
• Binary/shared-library mismatch: diagnose the reported package/path and compatibility.

Use one interpreter for pip and execution. Honor project requirements and compatible core/frontend revisions. Retaining installed PyTorch takes priority over setup examples. Optional CUDA/optimization warnings are not failures.

For isolation, create a **new unique** environment. From this skill's directory, use host pip 22.3+ and the bundled [hashed CPU wheels](scripts/torch-cpu-wheels.txt), then [requirements](scripts/requirements.txt):

set -e
cueq_python=python3  # Set to the selected project interpreter's path.
cueq_env=$(mktemp -d "${TMPDIR:-/tmp}/cueq-cpu.XXXXXX")
"$cueq_python" -m venv --without-pip "$cueq_env"
python3 -m pip --python "$cueq_env" install --cert "${PIP_CERT:-/etc/ssl/certs/ca-certificates.crt}" -q --no-deps -r scripts/torch-cpu-wheels.txt
python3 -m pip --python "$cueq_env" install --cert "${PIP_CERT:-/etc/ssl/certs/ca-certificates.crt}" -q -r scripts/requirements.txt
cueq_python="$cueq_env/bin/python"
"$cueq_python" -c 'import sys, torch, cuequivariance_torch; assert torch.version.cuda is None; print("CPU setup complete:", sys.executable, torch.__version__, cuequivariance_torch.__version__)'

Retain `cueq_python` for execution. For other platforms, use project setup. Without host pip, use installed `uv` with these requirements.

Use approved `PIP_CERT` for a private CA; retain TLS verification. If downloading fails, try a configured CPU mirror/cache or report the blocker. Do not replace CPU wheels with CUDA builds. Finish installs before retrying. See [pip's certificates](https://pip.pypa.io/en/stable/topics/https-certificates/).

Complete and inspect a run

Save the preamble and checks in `check.py`. Set `cueq_python` to the probe's `sys.executable`, or retain it from setup. Run **after setup exits successfully**:

set -e
cueq_run=$(mktemp -d "${TMPDIR:-/tmp}/cueq-check.XXXXXX")
printf 'run directory: %s\n' "$cueq_run"
if "$cueq_python" -u check.py >"$cueq_run/output.txt" 2>&1; then
    cueq_status=0
else
    cueq_status=$?
fi
printf '%s\n' "$cueq_status" >"$cueq_run/exit-code.txt"
tail -n 30 "$cueq_run/output.txt"
printf 'exit code: %s\n' "$cueq_status"
test "$cueq_status" -eq 0

The exit-code file appears only when Python finishes. If the tool yields, poll its session, then inspect the file and numerical output. Deliver the script, command, observed results, and backend explanation; label unrun checks. Keep install output separate.

Limitations and backends

| Method | Use and constraints | | --- | --- | | `naive` | Pure PyTorch reference implementation; supported CPU path and autograd. | | `uniform_1d` | CUDA; uniform segment shapes within each operand and a single mode after descriptor transformations. | | `fused_tp` | CUDA contractions with 3 or 4 operands; math dtype float32/float64. Some versions fall back on CPU, but select `naive` explicitly for a CPU reference. | | `indexed_linear` | CUDA indexed linear operations; indices must be sorted. |

High-level defaults vary: `SphericalHarmonics` selects `uniform_1