/bionemo-cuequivariance-jax

name: cuequivariance-jax
description: Use when writing or debugging cuEquivariance JAX polynomials, irrep dictionaries, Flax NNX layers, or CPU/CUDA execution and layouts.
license: Apache-2.0
metadata:
  author: Mario Geiger <mgeiger@nvidia.com>
  tags: [jax, equivariance, tensor-products]

cuEquivariance in JAX

Purpose

Execute `cuequivariance` (`cue`) descriptors with `cuequivariance_jax` (`cuex`). Preserve the project's framework and devices. Custom representations use core `cuequivariance`; PyTorch modules use `cuequivariance_torch`.

Inputs

Required inputs: operation, SO3/O3 irreps, layout, batch shapes, dtype, device, and weight ownership. Derive them from the request/code; state defaults. Graphs also need indices and destination count. Operand sizes/output irreps come from descriptors. Optional inputs: backend/performance preferences.

Instructions

1. Probe the project interpreter once and reuse working packages. **Environment setup/recovery must not recursively delete directories**, including temporary environments. Use a fresh unique path when isolation is needed. 2. Select `method="naive"` for CPU; Flax is needed only for NNX. High-level dictionary blocks are `(..., mul, ir.dim)`; uniform CUDA arrays differ. 3. Adapt one matching recipe into a script. Build descriptors outside transformed functions, broadcast leading input axes, and enable x64 **before array creation** for float64 checks. 4. Finish installation before running the saved script with the completion recipe below. A running session/task ID, empty output, or a file write is not a completed check. Poll the same session until it exits; then read its results before reporting success. 5. Investigate a concrete failure using its error and the affected installed API. Rerun affected checks after a fix. Once the requested checks pass, finish; avoid repeated probes, unrelated examples, or exploration of other backends for a CPU reference.

Prerequisites and environment recovery

Probe imports independently so a missing frontend does not hide working JAX. For an NNX task, include `"flax.nnx"` in the same loop:

import importlib
import sys

print("python:", sys.executable, sys.version.split()[0])
available, import_errors = {}, {}
for name in ("jax", "cuequivariance", "cuequivariance_jax"):
    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("import errors:", import_errors)
if "jax" in available:
    print("devices:", available["jax"].devices())

Reuse working packages; install only missing dependencies when needed. Honor the project's lockfile or supplied `scripts/requirements.txt` and compatible core/frontend revisions. JAX 0.11.1 / Flax 0.12.10 require **Python >=3.12**: select a compatible interpreter for those pins. The standalone combination below also supports Python 3.11; it does not replace project pins. Imports alone do not establish compatibility: cuEquivariance 0.13.0 with JAX 0.12.0 fails polynomial execution with a `segmented_polynomial_dce` argument-count error.

For isolation, keep the unique directory returned by `mktemp`; reuse it for reruns. This Linux recipe uses host pip 22.3+ and bypasses `ensurepip`:

set -e
cueq_python=python3  # Set to the selected project interpreter's path.
cueq_jax_env=$(mktemp -d "${TMPDIR:-/tmp}/cueq-jax.XXXXXX")
"$cueq_python" -m venv --without-pip "$cueq_jax_env"
python3 -m pip --python "$cueq_jax_env" install --cert "${PIP_CERT:-/etc/ssl/certs/ca-certificates.crt}" -q \
    "jax==0.10.2" "cuequivariance==0.13.0" "cuequivariance-jax==0.13.0"
cueq_python="$cueq_jax_env/bin/python"
JAX_PLATFORMS=cpu "$cueq_python" -c 'import jax, cuequivariance_jax; print("setup complete:", jax.__version__, cuequivariance_jax.__version__, jax.devices())'

Replace the package list with `-r` and the requirements path when prescribed; otherwise add `flax==0.12.7` to the same install command only for NNX. Keep this interpreter for execution; shell variables may need restoring in a new tool call. Without host pip, use an installed `uv`. The pip command uses the Linux system CA bundle; for a private CA use the environment's approved `PIP_CERT` bundle. Retain TLS verification. Diagnose the actual error before retrying; never reset directories or switch versions while an install is still running. If no compatible interpreter or trusted package source is available, report the blocker and unrun checks. See [pip's certificates](https://pip.pypa.io/en/stable/topics/https-certificates/).

Complete and inspect a run

Save the preamble and relevant checks in `check.py`. Set `cueq_python` to the probe's `sys.executable`, or retain it from setup. Run **after setup exits successfully**; a fresh directory records output and status:

set -e
cueq_run=$(mktemp -d "${TMPDIR:-/tmp}/cueq-check.XXXXXX")
printf 'run directory: %s\n' "$cueq_run"
if JAX_PLATFORMS=cpu "$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

If the tool yields, keep polling its session; the exit-code file is written only after Python finishes. Inspect that file and the numerical output. Empty output is not evidence that assertions ran. Deliver the script, actual command, and observed checks; label anything unrun. Keep setup output separate from the numerical results.

Limitations and backend selection

| Route | Supported use | | --- | --- | | `segmented_polynomial(..., method="naive")` | Portable JAX execution, including CPU, JIT, gradients, and vmap. | | `cuex.nnx.IrrepsLinear` | CPU/GPU channel mixing within matching irreps using JAX einsum. |