cuequivariance · diff
git:20260825.0e67a61 to git:20261002.53cded3
94 added, 381 removed. Audit A to A.
---
name: cuequivariance
- description: Define custom groups (Irrep subclasses), build segmented tensor products with CG coefficients, create equivariant polynomials and IrDictPolynomials, and use built-in descriptors (linear, tensor products, spherical harmonics). Use when working with cuequivariance group theory, irreps, or segmented polynomials.
+ 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: Groups, Irreps, and Segmented Polynomials
+ # cuEquivariance core
- ## Overview
+ 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.
- `cuequivariance` (imported as `cue`) provides two core abstractions:
+ ## Choose the workflow
- 1. **Group theory**: `Irrep` subclasses define irreducible representations of Lie groups (SO3, O3, SU2, or custom). `Irreps` manages collections with multiplicities.
- 2. **Segmented polynomials**: `SegmentedTensorProduct` describes tensor contractions over segments of varying shape, linked by `Path` objects carrying Clebsch-Gordan coefficients. `SegmentedPolynomial` wraps multiple STPs into a polynomial with named inputs/outputs. Two higher-level wrappers attach group representations:
- - `EquivariantPolynomial` — dense operands with `IrrepsAndLayout` metadata
- - `IrDictPolynomial` — operands already split by irrep, with per-group `Irreps` metadata for the `dict[Irrep, Array]` 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.
- ## Defining a custom group
+ ## Check the installed API
- Subclass `cue.Irrep` (a frozen dataclass) and implement:
+ These examples were checked with **cuEquivariance 0.12.0**. Check the user's
+ environment before using version-sensitive APIs, especially `IrDictPolynomial`
+ and the `_ir_dict` descriptors:
```python
- from __future__ import annotations
- import dataclasses
- import re
- from typing import Iterator
- import numpy as np
+ from importlib.metadata import version
import cuequivariance as cue
- @dataclasses.dataclass(frozen=True)
- class Z2(cue.Irrep):
- odd: bool # dataclass field -- required for correct __eq__ and __hash__
-
- # No __init__ needed -- @dataclass(frozen=True) generates it: Z2(odd=True)
-
- @classmethod
- def regexp_pattern(cls) -> re.Pattern:
- # Pattern whose first group is passed to from_string
- return re.compile(r"(odd|even)")
-
- @classmethod
- def from_string(cls, string: str) -> Z2:
- return cls(odd=string == "odd")
-
- def __repr__(rep: Z2) -> str:
- return "odd" if rep.odd else "even"
-
- def __mul__(rep1: Z2, rep2: Z2) -> Iterator[Z2]:
- # Selection rule: which irreps appear in the tensor product rep1 x rep2
- return [Z2(odd=rep1.odd ^ rep2.odd)]
-
- @classmethod
- def clebsch_gordan(cls, rep1: Z2, rep2: Z2, rep3: Z2) -> np.ndarray:
- # Shape: (num_paths, rep1.dim, rep2.dim, rep3.dim)
- if rep3 in rep1 * rep2:
- return np.array([[[[1]]]])
- else:
- return np.zeros((0, 1, 1, 1))
-
- @property
- def dim(rep: Z2) -> int:
- return 1
-
- def __lt__(rep1: Z2, rep2: Z2) -> bool:
- # Ordering for sorting; dimension is compared first by the base class
- return rep1.odd < rep2.odd
-
- @classmethod
- def iterator(cls) -> Iterator[Z2]:
- # Must yield trivial irrep first
- for odd in [False, True]:
- yield Z2(odd=odd)
-
- def discrete_generators(rep: Z2) -> np.ndarray:
- # Shape: (num_generators, dim, dim)
- if rep.odd:
- return -np.ones((1, 1, 1))
- else:
- return np.ones((1, 1, 1))
-
- def continuous_generators(rep: Z2) -> np.ndarray:
- # Shape: (lie_dim, dim, dim) -- Z2 is discrete, so lie_dim=0
- return np.zeros((0, rep.dim, rep.dim))
-
- def algebra(self) -> np.ndarray:
- # Shape: (lie_dim, lie_dim, lie_dim) -- structure constants [X_i, X_j] = A_ijk X_k
- return np.zeros((0, 0, 0))
-
-
- # Usage:
- irreps = cue.Irreps(Z2, "3x odd + 2x even") # dim=5
- ```
-
- ### Required methods summary
-
- | Method | Returns | Purpose |
- |--------|---------|---------|
- | `regexp_pattern()` | `re.Pattern` | Parse string like `"1"`, `"0e"`, `"odd"` |
- | `from_string(s)` | `Irrep` | Construct from matched string |
- | `__repr__` | `str` | Canonical string form |
- | `__mul__(a, b)` | `Iterator[Irrep]` | Selection rule for tensor product |
- | `clebsch_gordan(a, b, c)` | `ndarray (n, d1, d2, d3)` | CG coefficients |
- | `dim` (property) | `int` | Dimension of representation |
- | `__lt__(a, b)` | `bool` | Ordering (dimension first, then custom) |
- | `iterator()` | `Iterator[Irrep]` | All irreps, trivial first |
- | `continuous_generators()` | `ndarray (lie_dim, dim, dim)` | Lie algebra generators |
- | `discrete_generators()` | `ndarray (n, dim, dim)` | Finite symmetry generators |
- | `algebra()` | `ndarray (lie_dim, lie_dim, lie_dim)` | Structure constants |
-
- ### Built-in groups
-
- - **`cue.SO3(l)`**: 3D rotations. `l` is a non-negative integer. `dim = 2l+1`. String: `"0"`, `"1"`, `"2"`.
- - **`cue.O3(l, p)`**: 3D rotations + parity. `p=+1` (even) or `p=-1` (odd). String: `"0e"`, `"1o"`, `"2e"`.
- - **`cue.SU2(j)`**: Spin group. `j` is a non-negative half-integer. String: `"0"`, `"1/2"`, `"1"`.
-
- ## Irreps and layout
-
- ```python
- irreps = cue.Irreps("SO3", "16x0 + 4x1 + 2x2") # 16 scalars, 4 vectors, 2 rank-2
- irreps.dim # 16*1 + 4*3 + 2*5 = 38
-
- for mul, ir in irreps:
- print(mul, ir, ir.dim) # 16 0 1, then 4 1 3, then 2 2 5
- ```
-
- `IrrepsLayout` controls memory order within each `(mul, ir)` block:
-
- - `cue.ir_mul`: data ordered as `(ir.dim, mul)` — **used by all descriptors and ir_dict internally**
- - `cue.mul_ir`: data ordered as `(mul, ir.dim)` — **used by nnx `dict[Irrep, Array]` and PyTorch**
-
- `IrrepsAndLayout` combines irreps with a layout into a `Rep`:
-
- ```python
- rep = cue.IrrepsAndLayout(cue.Irreps("SO3", "4x0 + 2x1"), cue.ir_mul)
- rep.dim # 10
- ```
-
- ## Building a SegmentedTensorProduct from scratch
-
- The subscripts string uses Einstein notation. Operands are comma-separated, coefficient modes follow `+`.
-
- ```python
- # Matrix-vector multiply: y_i = sum_j M_ij * x_j
- d = cue.SegmentedTensorProduct.from_subscripts("ij,j,i")
- d.add_segment(0, (3, 4)) # operand 0: matrix segment of shape (3, 4)
- d.add_segment(1, (4,)) # operand 1: vector of size 4
- d.add_segment(2, (3,)) # operand 2: output vector of size 3
- d.add_path(0, 0, 0, c=1.0) # link segments 0,0,0 with coefficient=1.0
-
- poly = cue.SegmentedPolynomial.eval_last_operand(d) # last operand becomes output
- [y] = poly(M_flat, x) # numpy evaluation
- ```
-
- ### Multi-segment STP (how descriptors work internally)
-
- Descriptors build STPs with multiple segments per operand. Each segment corresponds to an irrep block:
-
- ```python
- # Linear equivariant map: output[iv] = sum_u weight[uv] * input[iu]
- d = cue.SegmentedTensorProduct.from_subscripts("uv,iu,iv")
-
- # Segment for l=1: ir_dim=3, mul_in=2, mul_out=5
- s_in_0 = d.add_segment(1, (3, 2)) # input block
- s_out_0 = d.add_segment(2, (3, 5)) # output block
- d.add_path((2, 5), s_in_0, s_out_0, c=1.0)
-
- # Segment for l=0: ir_dim=1, mul_in=4, mul_out=3
- s_in_1 = d.add_segment(1, (1, 4))
- s_out_1 = d.add_segment(2, (1, 3))
- d.add_path((4, 3), s_in_1, s_out_1, c=1.0)
- ```
-
- ### Weights operand
-
- For weighted tensor products (subscript starting with `uvw` or `uv`), the first operand is always weights. The weight segment shape is `(mul_1, mul_2, ...)` matching the multiplicity modes. The weights operand gets `new_scalars()` irreps since weights are invariant.
-
- ### CG coefficients as path coefficients
-
- ```python
- d = cue.SegmentedTensorProduct.from_subscripts("uvw,iu,jv,kw+ijk")
- # For each pair of input irreps and each output irrep in the selection rule:
- for cg in cue.clebsch_gordan(ir1, ir2, ir3):
- # cg has shape (ir1.dim, ir2.dim, ir3.dim)
- d.add_path((mul1, mul2, mul3), seg_in1, seg_in2, seg_out, c=cg)
- ```
-
- ## Descriptors
-
- All descriptors come in two variants:
-
- - **Original** — returns `EquivariantPolynomial` with dense operands
- - **`_ir_dict`** — returns `IrDictPolynomial` with operands already split by irrep
-
- ### EquivariantPolynomial descriptors
-
- ```python
- # Fully connected tensor product (all input-output irrep combinations)
- e = cue.descriptors.fully_connected_tensor_product(
- 16 * cue.Irreps("SO3", "0 + 1 + 2"),
- 16 * cue.Irreps("SO3", "0 + 1 + 2"),
- 16 * cue.Irreps("SO3", "0 + 1 + 2"),
- )
-
- # Channelwise tensor product (same-channel only, sparse)
- e = cue.descriptors.channelwise_tensor_product(
- 64 * cue.Irreps("SO3", "0 + 1"), cue.Irreps("SO3", "0 + 1"),
- cue.Irreps("SO3", "0 + 1"), simplify_irreps3=True,
- )
-
- # Full (weightless) tensor product
- e = cue.descriptors.full_tensor_product(
- cue.Irreps("SO3", "2x0 + 1x1"), cue.Irreps("SO3", "0 + 1"),
- )
-
- # Elementwise tensor product (paired channels)
- e = cue.descriptors.elementwise_tensor_product(
- cue.Irreps("SO3", "4x0 + 4x1"), cue.Irreps("SO3", "4x0 + 4x1"),
- )
-
- # Linear equivariant map (weight x input)
- e = cue.descriptors.linear(
- cue.Irreps("SO3", "4x0 + 2x1"),
- cue.Irreps("SO3", "3x0 + 5x1"),
- )
-
- # Spherical harmonics
- e = cue.descriptors.spherical_harmonics(cue.SO3(1), [0, 1, 2, 3])
-
- # Symmetric contraction (MACE-style)
- e = cue.descriptors.symmetric_contraction(
- 64 * cue.Irreps("SO3", "0 + 1 + 2"),
- 64 * cue.Irreps("SO3", "0 + 1"),
- (1, 2, 3),
- )
- ```
-
- ### IrDictPolynomial descriptors
-
- Each `_ir_dict` variant returns an `IrDictPolynomial` whose polynomial is already split by irrep. The `input_irreps` and `output_irreps` tuples describe the operand groups.
-
- ```python
- # Channelwise tensor product
- desc = cue.descriptors.channelwise_tensor_product_ir_dict(
- 64 * cue.Irreps("SO3", "0 + 1"),
- cue.Irreps("SO3", "0 + 1"),
- cue.Irreps("SO3", "0 + 1"),
- )
- # desc.polynomial — SegmentedPolynomial, already split by irrep
- # desc.input_irreps — (weight_irreps, irreps1, irreps2)
- # desc.output_irreps — (irreps_out,)
-
- # Fully connected tensor product
- desc = cue.descriptors.fully_connected_tensor_product_ir_dict(irreps1, irreps2, irreps3)
-
- # Full (weightless) tensor product
- desc = cue.descriptors.full_tensor_product_ir_dict(irreps1, irreps2)
-
- # Elementwise tensor product
- desc = cue.descriptors.elementwise_tensor_product_ir_dict(irreps1, irreps2)
-
- # Linear
- desc = cue.descriptors.linear_ir_dict(irreps_in, irreps_out)
-
- # Spherical harmonics
- desc = cue.descriptors.spherical_harmonics_ir_dict(cue.O3(1, -1), [0, 1, 2, 3])
-
- # Symmetric contraction
- desc = cue.descriptors.symmetric_contraction_ir_dict(irreps_in, irreps_out, (1, 2, 3))
- ```
-
- ### IrDictPolynomial
-
- `IrDictPolynomial` pairs a `SegmentedPolynomial` (already split by irrep) with the `Irreps` that describe each operand group.
-
- ```python
- desc = cue.descriptors.channelwise_tensor_product_ir_dict(
- 32 * cue.Irreps("SO3", "0 + 1"),
- cue.Irreps("SO3", "0 + 1"),
- cue.Irreps("SO3", "0 + 1"),
- )
-
- desc.polynomial # SegmentedPolynomial — each operand is one (mul, ir) block
- desc.input_irreps # (weight_irreps, irreps1, irreps2)
- desc.output_irreps # (irreps_out,)
-
- # Scale coefficients
- scaled_poly = desc.polynomial * 0.5
-
- # Access individual operand info
- for i, op in enumerate(desc.polynomial.inputs):
- print(f"Input {i}: size={op.size}, num_segments={op.num_segments}")
- ```
-
- Contract: for each `(mul, ir)` block in `input_irreps` / `output_irreps`, the corresponding polynomial operand has size `mul * ir.dim`.
-
- ### split_polynomial_by_irreps
-
- The low-level function underlying `_ir_dict` descriptors. Splits one polynomial operand at irrep boundaries:
-
- ```python
- poly = e.polynomial # from an EquivariantPolynomial
- poly = cue.split_polynomial_by_irreps(poly, 2, irreps_sh) # split input 2
- poly = cue.split_polynomial_by_irreps(poly, 1, irreps_in) # split input 1
- poly = cue.split_polynomial_by_irreps(poly, -1, irreps_out) # split output
+ print(version("cuequivariance"))
+ print("IrDictPolynomial available:", hasattr(cue, "IrDictPolynomial"))
```
- ### EquivariantPolynomial key methods
-
- ```python
- e.inputs # tuple of Rep (group representations for each input)
- e.outputs # tuple of Rep
- e.polynomial # the underlying SegmentedPolynomial
-
- # Numpy evaluation
- [out] = e(weights, input1, input2)
-
- # Preparing for uniform_1d execution (see cuequivariance_jax SKILL.md)
- e_ready = e.squeeze_modes().flatten_coefficient_modes()
-
- # Split an operand into per-irrep pieces (for ir_dict interface)
- e_split = e.split_operand_by_irrep(1).split_operand_by_irrep(-1)
-
- # Scale all coefficients
- e_scaled = e * 0.5
+ If the package is absent, install it in the project's environment, preserving
+ existing dependency constraints. A standalone CPU example environment can use
+ `python -m pip install "cuequivariance==0.12.0"`. Core examples need NumPy but
+ neither CUDA nor JAX/PyTorch. If an API is unavailable, use the supported dense
+ descriptor or agree on a compatible upgrade; do not invent a substitute method.
- # Fuse compatible STPs
- e_fused = e.fuse_stps()
- ```
+ 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.
- ### normalize_paths_for_operand
+ ## Groups, multiplicities, and layout
- Called internally by descriptors. Normalizes path coefficients so that a random input produces unit-variance output for the specified operand. Critical for numerical stability.
+ | 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 |
- ## SegmentedPolynomial structure
+ `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`.
- ```python
- poly = e.polynomial
- poly.num_inputs # number of input operands
- poly.num_outputs # number of output operands
- poly.inputs # tuple of SegmentedOperand
- poly.outputs # tuple of SegmentedOperand
- poly.operations # tuple of (Operation, SegmentedTensorProduct)
+ Layout is defined **within each irrep block**, not across the whole vector:
- # Each operation maps buffers to STP operands
- for op, stp in poly.operations:
- print(op.buffers) # e.g., (0, 1, 2) means inputs[0], inputs[1] -> outputs[0]
- print(stp.subscripts)
- ```
+ | Layout | Block shape before flattening |
+ | --- | --- |
+ | `cue.ir_mul` | `(..., ir.dim, mul)` |
+ | `cue.mul_ir` | `(..., mul, ir.dim)` |
- ### SegmentedOperand
+ 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.
- ```python
- operand = poly.inputs[0]
- operand.num_segments # how many segments
- operand.segments # tuple of shape tuples, e.g., ((3, 4), (1, 2))
- operand.size # total flattened size (sum of products of segment shapes)
- operand.ndim # number of dimensions per segment
- operand.all_same_segment_shape() # True if all segments have identical shape
- operand.segment_shape # the common shape (only if all_same_segment_shape)
- ```
+ ## Runnable dense tensor product
- ## Custom equivariant polynomial from scratch
+ 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
- # Build a fully-connected SO3(1)xSO3(1)->SO3(0) tensor product manually
- cg = cue.clebsch_gordan(cue.SO3(1), cue.SO3(1), cue.SO3(0)) # shape (1, 3, 3, 1)
+ 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()
- d = cue.SegmentedTensorProduct.from_subscripts("uvw,iu,jv,kw+ijk")
- d.add_segment(1, (3, 4)) # input1: 4x SO3(1), shape=(ir_dim, mul)
- d.add_segment(2, (3, 4)) # input2: 4x SO3(1)
- d.add_segment(3, (1, 16)) # output: 16x SO3(0) (4*4 fully connected)
+ # 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")
+ ```
- for c in cg:
- d.add_path((4, 4, 16), 0, 0, 0, c=c)
+ ## Select a descriptor
- d = d.normalize_paths_for_operand(-1)
+ All names below are under `cue.descriptors`. Inspect the returned input/output
+ metadata to allocate arrays, including weights; do not guess the weight count.
- poly = cue.SegmentedPolynomial.eval_last_operand(d)
- ep = cue.EquivariantPolynomial(
- [
- cue.IrrepsAndLayout(cue.Irreps("SO3", "4x1").new_scalars(d.operands[0].size), cue.ir_mul),
- cue.IrrepsAndLayout(cue.Irreps("SO3", "4x1"), cue.ir_mul),
- cue.IrrepsAndLayout(cue.Irreps("SO3", "4x1"), cue.ir_mul),
- ],
- [cue.IrrepsAndLayout(cue.Irreps("SO3", "16x0"), cue.ir_mul)],
- poly,
- )
+ | 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 |
- # Numpy evaluation
- w = np.random.randn(ep.inputs[0].dim)
- x = np.random.randn(ep.inputs[1].dim)
- y = np.random.randn(ep.inputs[2].dim)
- [out] = ep(w, x, y)
- ```
+ 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.
- ## Key file locations
+ `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.
- | Component | Path |
- |-----------|------|
- | `Irrep` base class | `cuequivariance/group_theory/representations/irrep.py` |
- | `Rep` base class | `cuequivariance/group_theory/representations/rep.py` |
- | `SO3` | `cuequivariance/group_theory/representations/irrep_so3.py` |
- | `O3` | `cuequivariance/group_theory/representations/irrep_o3.py` |
- | `SU2` | `cuequivariance/group_theory/representations/irrep_su2.py` |
- | `Irreps` | `cuequivariance/group_theory/irreps_array/irreps.py` |
- | `IrrepsLayout` | `cuequivariance/group_theory/irreps_array/irreps_layout.py` |
- | `IrrepsAndLayout` | `cuequivariance/group_theory/irreps_array/irreps_and_layout.py` |
- | `SegmentedTensorProduct` | `cuequivariance/segmented_polynomials/segmented_tensor_product.py` |
- | `SegmentedPolynomial` | `cuequivariance/segmented_polynomials/segmented_polynomial.py` |
- | `EquivariantPolynomial` | `cuequivariance/group_theory/equivariant_polynomial.py` |
- | `IrDictPolynomial` | `cuequivariance/group_theory/ir_dict_polynomial.py` |
- | Descriptors | `cuequivariance/group_theory/descriptors/` |
- | Tensor product descriptors | `cuequivariance/group_theory/descriptors/irreps_tp.py` |
- | `spherical_harmonics` | `cuequivariance/group_theory/descriptors/spherical_harmonics_.py` |
- | `symmetric_contraction` | `cuequivariance/group_theory/descriptors/symmetric_contractions.py` |
+ 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.