cuequivariance git:20261002.b8a760bA

Immutable. This exact content never changes and is served at /api/v1/blob/3e5135ba4ef7bddc.

---
name: cuequivariance
description: Build and debug cuEquivariance irreps, custom Irrep subclasses, Clebsch-Gordan tensor products, and equivariant or segmented polynomials. Use for cuequivariance core descriptors and IrDictPolynomial layout or shape questions; framework execution requires the matching JAX or PyTorch API.
---

# cuEquivariance core

Use `cuequivariance` (`import cuequivariance as cue`) to describe group
representations and tensor contractions. Prefer a built-in descriptor when it
expresses the requested operation; construct a `SegmentedTensorProduct` only
for a custom contraction. Core polynomials can be evaluated with NumPy on CPU.
They do not create trainable framework parameters or select a GPU kernel.

## Choose the workflow

- For built-in groups and dense descriptors, use the examples below.
- For a new group, read [custom irreps](references/custom-irreps.md).
- For custom contractions, CG paths, or polynomial inspection, read
  [segmented polynomials](references/segmented-polynomials.md).
- For dictionaries keyed by irrep, read [the ir_dict workflow](references/ir-dict.md).
  It includes a CPU example and the boundary with JAX execution.

For explanations or requests to show runnable code, adapt the closest example;
installing packages is not a prerequisite to answering. Read the reference for
the selected workflow. Execute code when the user requests a run or when an
integration or correctness question needs a runtime check. Inspect upstream
source when the installed API differs or a numerical check fails.

## Check the execution environment

These examples were checked with **cuEquivariance 0.12.0**. When executing,
check the selected interpreter and the APIs needed for that workflow together:

```python
from importlib.metadata import PackageNotFoundError, version
import sys

print("Python:", sys.executable)
try:
    installed = version("cuequivariance")
except PackageNotFoundError:
    print("cuequivariance is not installed in this interpreter")
else:
    import cuequivariance as cue

    print("cuEquivariance:", installed)
    print("IrDictPolynomial available:", hasattr(cue, "IrDictPolynomial"))
```

Reuse a compatible project environment and its dependency constraints. If
execution needs missing packages, use [isolated example execution](references/environment.md).
Preserve existing virtual environments and the system/user Python installation;
do not delete environments or bypass package-manager protections to repair setup.
When installation is unavailable or disallowed, provide the runnable code and
state which checks were not executed. Core examples need NumPy but neither CUDA
nor JAX/PyTorch. If an API is unavailable, use a supported dense descriptor when
it satisfies the task; do not invent a method or force an environment upgrade.

For a framework task, inspect the installed `cuequivariance_jax` or
`cuequivariance_torch` signature and its version's backend requirements.
[Official documentation](https://docs.nvidia.com/cuda/cuequivariance/) and the
[0.12.0 source](https://github.com/NVIDIA/cuEquivariance/tree/v0.12.0) cover those
separate packages.

## Groups, multiplicities, and layout

| Group | Construction | Dimension and meaning |
| --- | --- | --- |
| SO(3) | `cue.SO3(l)` | `2*l + 1`; rotations; nonnegative integer `l` |
| O(3) | `cue.O3(l, p)` | `2*l + 1`; also parity `p=+1` or `-1`, e.g. `0e`, `1o` |
| SU(2) | `cue.SU2(j)` | `2*j + 1`; nonnegative integer or half-integer spin |

`cue.Irreps("SO3", "16x0 + 4x1 + 2x2")` has dimension 38. Iterating it yields
`(multiplicity, irrep)` blocks. Multiplying an `Irreps` by an integer scales
multiplicities; `ir1 * ir2` on individual irreps gives the selection rule.
For SO(3), allowed output angular momenta run from `abs(l1-l2)` to `l1+l2`.
For O(3), output parity must also equal `p1*p2`.

Layout is defined **within each irrep block**, not across the whole vector:

| Layout | Block shape before flattening |
| --- | --- |
| `cue.ir_mul` | `(..., ir.dim, mul)` |
| `cue.mul_ir` | `(..., mul, ir.dim)` |

Read `e.inputs[i].layout` and `e.outputs[i].layout` from the descriptor.
The tensor-product examples here use `ir_mul`; many PyTorch layers default to
`mul_ir` and accept explicit layout options. JAX dictionary features often use
`mul_ir`, while polynomial execution can require `ir_mul`.
Convert the actual data block by block with a transpose before flattening;
changing metadata or calling `reshape` alone does not change the element order.
Similarly, sorting or simplifying irreps metadata does not reorder an existing
tensor. Preserve the descriptor's block order when evaluating it.

## Runnable dense tensor product

When the user asks for all output irreps, start with `full_tensor_product`.
It has two inputs and no weights. This example includes all input definitions:

```python
import numpy as np
import cuequivariance as cue

irreps1 = cue.Irreps("SO3", "2x0 + 1x1 + 1x2")
irreps2 = cue.Irreps("SO3", "1x0 + 1x1")
e = cue.descriptors.full_tensor_product(irreps1, irreps2)
rng = np.random.default_rng(0)
x1 = rng.normal(size=(4, e.inputs[0].dim))
x2 = rng.normal(size=(4, e.inputs[1].dim))
[out] = e(x1, x2)
assert out.shape == (4, 40)
assert np.isfinite(out).all()

# A multiplicity summary for display; keep e.outputs[0] for out's actual layout.
summary = e.outputs[0].irreps.sort().irreps.simplify()
assert summary == cue.Irreps("SO3", "3x0 + 5x1 + 3x2 + 1x3")
```

## Select a descriptor

All names below are under `cue.descriptors`. Inspect the returned input/output
metadata to allocate arrays, including weights; do not guess the weight count.

| Descriptor | Use and operand order |
| --- | --- |
| `linear(irreps_in, irreps_out)` | Mix multiplicities of matching irreps; weights, input → output |
| `fully_connected_tensor_product(a, b, out)` | Choose exact output multiplicities; weights, input1, input2 → output |
| `full_tensor_product(a, b)` | All allowed couplings, no weights; input1, input2 → output |
| `channelwise_tensor_product(a, b, irreps3_filter=...)` | Weighted channel pairs; output multiplicity per path is `mul1*mul2` |
| `elementwise_tensor_product(a, b)` | Paired channels without weights; inputs need equal numbers of irrep copies and compatible block boundaries |
| `spherical_harmonics(cue.O3(1, -1), [0, 1, 2, 3])` | Vector → harmonics with parity; normalize input vectors separately if unit-sphere values are required |
| `symmetric_contraction(irreps_in, irreps_out, (1, 2, 3))` | MACE-style polynomial correlations; inspect descriptor shapes and multiplicity constraints |

The third argument to a channelwise product **filters irrep types**; it does not
set output multiplicities. Channelwise retains pairs of channels (`u, v`), unlike
elementwise's matched channel (`u`). Use `simplify_irreps3=True` when a dense
channelwise descriptor should consolidate repeated output irreps.

`EquivariantPolynomial` stores dense representation metadata in `inputs` and
`outputs`, and its contraction in `polynomial`. Calling it returns a **list** of
NumPy outputs, even for one output. The corresponding `_ir_dict` variants return
metadata plus a polynomial already split by irrep; see the linked workflow.

For a custom implementation, check output dimensions and finite values, then
verify covariance under a nontrivial group transformation. Use the same group,
basis, block order, and layout for both the inputs and expected output.