cuequivariance

What problem does it solve?

cuEquivariance provides a framework to define custom groups, construct segmented tensor products with Clebsch-Gordan coefficients, and compose equivariant polynomials using built-in descriptors such as linear, fully_connected_tensor_product, channelwise_tensor_product, and spherical harmonics. This enables researchers and developers to prototype and deploy equivariant neural networks that respect symmetries in 3D data.

Core Features & Use Cases

• Define custom groups by subclassing cue.Irrep and implementing required methods such as regexp_pattern, from_string, repr, mul, clebsch_gordan, and dimension, enabling custom symmetry handling.
• Build SegmentedTensorProduct and SegmentedPolynomial from arbitrary subscripts to describe segmented contractions linked by Path objects carrying CG coefficients.
• Use higher-level descriptors (fully_connected_tensor_product, channelwise_tensor_product, elementwise_tensor_product, linear, spherical_harmonics, symmetric_contraction) to construct EquivariantPolynomial objects for neural networks with SO3/O3/SU2 groups or custom groups.
• Evaluate and manipulate EquivariantPolynomial and SegmentedPolynomial using a variety of utilities (stacking, fusing, squeezing, canonicalization, and automatic differentiation helpers).

Quick Start

Build a tiny SO3 descriptor with a linear map and evaluate it on dummy inputs.