cuequivariance

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.
• For custom contractions, CG paths, or polynomial inspection, read segmented polynomials.
• For dictionaries keyed by irrep, read the ir_dict workflow. 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:

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. 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 and the 0.12.0 source 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:

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.