Execute equivariant polynomials in JAX using segmented_polynomial (naive/uniform_1d), the ir_dict workflow with IrDictPolynomial and dict[Irrep, Array], and Flax NNX layers (IrrepsLinear, SphericalHarmonics, IrrepsIndexedLinear). Use when writing JAX code with cuequivariance.
Execute equivariant polynomials in JAX using segmented_polynomial (naive/uniform_1d), the ir_dict workflow with IrDictPolynomial and dict[Irrep, Array], and Flax NNX layers (IrrepsLinear, SphericalHarmonics, IrrepsIndexedLinear). Use when writing JAX code with cuequivariance.
cuequivariance_jax: Executing Equivariant Polynomials in JAX
Overview
cuequivariance_jax (imported as cuex) executes cuequivariance polynomials on GPU via JAX. It provides:
Core primitive: cuex.segmented_polynomial() — JAX primitive with full AD/vmap/JIT support
Two data representations (both built on segmented_polynomial):
cuex.equivariant_polynomial() + RepArray — the original interface, a single contiguous array with representation metadata
cuex.ir_dict module — dict[Irrep, Array] interface, uses IrDictPolynomial descriptors, works naturally with jax.tree
GPU, all segments uniform shape within each operand, single mode
"indexed_linear"
CUDA kernel
GPU, linear operations with cuex.Repeats indexing
Core primitive: segmented_polynomial
import jax
import jax.numpy as jnp
import cuequivariance as cue
import cuequivariance_jax as cuex
# Build a descriptor
e = cue.descriptors.channelwise_tensor_product(
32 * cue.Irreps("SO3", "0 + 1"),
cue.Irreps("SO3", "0 + 1"),
cue.Irreps("SO3", "0 + 1"),
)
poly = e.polynomial
batch = 64
w = jnp.ones((poly.inputs[0].size,)) # weights (shared across batch)
x = jax.random.normal(key, (batch, poly.inputs[1].size)) # batched input 1
y = jax.random.normal(key, (batch, poly.inputs[2].size)) # batched input 2# Execute with naive method
[out] = cuex.segmented_polynomial(
poly,
[w, x, y], # inputs
[jax.ShapeDtypeStruct((batch, poly.outputs[0].size), jnp.float32)], # output spec
method="naive",
)
# Execute with uniform_1d (GPU, requires uniform segments)
[out] = cuex.segmented_polynomial(
poly, [w, x, y],
[jax.ShapeDtypeStruct((batch, poly.outputs[0].size), jnp.float32)],
method="uniform_1d",
)
Multiple batch axes with broadcasting
Inputs can have any number of batch axes (everything before the last axis). Standard NumPy broadcasting applies: each batch axis is either size-1 or a common size. Inputs with fewer batch dimensions are implicitly prepended with size-1 axes:
# Fewer batch dims: weights with no batch axis broadcast across all
w = jnp.ones((poly.inputs[0].size,)) # 0 batch axes -> broadcasts
x = jnp.ones((5, 10, poly.inputs[1].size))
y = jnp.ones((5, 10, poly.inputs[2].size))
[out] = cuex.segmented_polynomial(
poly, [w, x, y],
[jax.ShapeDtypeStruct((5, 10, poly.outputs[0].size), jnp.float32)],
method="uniform_1d",
)
Indexing (gather/scatter)
Index arrays provide gather (for inputs) and scatter (for outputs). One index per operand (inputs + outputs), None means no indexing:
Uses dict[Irrep, Array] where each value has shape (..., multiplicity, irrep_dim). This is the standard representation for NNX layers and works naturally with jax.tree operations.
Getting an ir_dict-ready polynomial
Use _ir_dict descriptor variants, which return IrDictPolynomial with the polynomial already split by irrep:
desc = cue.descriptors.channelwise_tensor_product_ir_dict(
32 * cue.Irreps("SO3", "0 + 1"),
cue.Irreps("SO3", "0 + 1"),
cue.Irreps("SO3", "0 + 1"),
)
poly = desc.polynomial # SegmentedPolynomial, already split by irrep
weight_irreps, irreps1, irreps2 = desc.input_irreps
(irreps_out,) = desc.output_irreps # tuple unpacking to get the single output group
Each polynomial operand corresponds to exactly one (mul, ir) block. The input_irreps and output_irreps tuples describe how operands group into logical operand groups (weights, node features, spherical harmonics, output).
Executing with segmented_polynomial_uniform_1d
from einops import rearrange
num_edges, num_nodes = 100, 30# Weights: reshape to (batch, num_segments, segment_size)
w_flat = jax.random.normal(key, (num_edges, poly.inputs[0].size))
w = rearrange(w_flat, "e (s m) -> e s m", s=poly.inputs[0].num_segments)
# Node features: dict[Irrep, Array] reshaped to (nodes, ir.dim, mul) for ir_mul layout
node_feats = {
cue.SO3(0): jnp.ones((num_nodes, 32, 1)), # 32x scalar
cue.SO3(1): jnp.ones((num_nodes, 32, 3)), # 32x vector
}
x1 = jax.tree.map(lambda v: rearrange(v, "n m i -> n i m"), node_feats)
# Spherical harmonics: (edges, ir.dim) — no multiplicity dimension
sph = {
cue.SO3(0): jnp.ones((num_edges, 1)),
cue.SO3(1): jnp.ones((num_edges, 3)),
}
# Build output template
senders = jax.random.randint(key, (num_edges,), 0, num_nodes)
receivers = jax.random.randint(key, (num_edges,), 0, num_nodes)
out_template = {
ir: jax.ShapeDtypeStruct(
(num_nodes, desc.num_segments) + desc.segment_shape, w.dtype
)
for (_, ir), desc inzip(irreps_out, poly.outputs)
}
# Execute with gather (senders) and scatter (receivers)
y = cuex.ir_dict.segmented_polynomial_uniform_1d(
poly,
[w, x1, sph],
out_template,
input_indices=[None, senders, None],
output_indices=receivers,
name="tensor_product",
)
# y is dict[Irrep, Array] with accumulated results at receiver nodes
norm = cuex.nnx.IrrepsNormalize(eps=1e-6, scale=1.0, skip_scalars=True)
y = norm(x) # normalizes non-scalar irreps by RMS over ir.dim, averaged over mul
MLP (scalar only)
mlp = cuex.nnx.MLP(
layer_sizes=[64, 128, 64],
activation=jax.nn.silu,
output_activation=False,
dtype=jnp.float32,
rngs=nnx.Rngs(0),
)
y = mlp(x_scalar) # standard dense MLP with 1/sqrt(fan_in) normalization
IrrepsIndexedLinear
For species-indexed linear layers (different weights per atom type):
indexed_linear = cuex.nnx.IrrepsIndexedLinear(
irreps_in=cue.Irreps(cue.O3, "8x0e").regroup(),
irreps_out=cue.Irreps(cue.O3, "16x0e").regroup(),
num_indices=50, # number of species
scale=1.0,
dtype=jnp.float32,
rngs=nnx.Rngs(0),
)
# num_index_counts: how many atoms of each species
species_counts = jnp.array([3, 4, 3, ...]) # sum = batch_size
y = indexed_linear(x, species_counts)
Uses method="indexed_linear" internally with cuex.Repeats.
Preparing polynomials for uniform_1d
The uniform_1d CUDA kernel requires:
All segments within each operand have the same shape
A single mode in the subscripts (after preprocessing)
From EquivariantPolynomial to uniform_1d-ready
For equivariant_polynomial() (RepArray interface):
e = cue.descriptors.channelwise_tensor_product(...)
e = e.squeeze_modes().flatten_coefficient_modes()
out = cuex.equivariant_polynomial(e, inputs, method="uniform_1d")
For ir_dict (dict[Irrep, Array] interface), use _ir_dict descriptors directly:
desc = cue.descriptors.channelwise_tensor_product_ir_dict(
irreps_in, irreps_sh, irreps_out
)
poly = desc.polynomial
# Each operand has a single irrep type -> maps naturally to dict[Irrep, Array]
Why splitting by irrep matters
Without splitting, a dense operand like 32x0+32x1 requires all irreps packed into a single contiguous buffer. After splitting, each irrep gets its own separate buffer passed to the CUDA kernel via FFI. The buffers no longer need to be contiguous with each other.
This is especially useful when the polynomial is preceded or followed by per-irrep linear layers (like IrrepsLinear). With split operands, no transpose or copy is needed between the linear layers and the polynomial — the dict[Irrep, Array] flows directly through the pipeline.
Complete GNN message-passing example
This pattern is used in NequIP, MACE, and similar equivariant GNN models:
classMessagePassing(nnx.Module):
def__init__(self, irreps_in, irreps_sh, irreps_out, epsilon, *, name, dtype, rngs):
self.name = name
desc = cue.descriptors.channelwise_tensor_product_ir_dict(
irreps_in, irreps_sh, irreps_out
)
(self.irreps_out,) = desc.output_irreps
self.poly = desc.polynomial * epsilon
self.weight_numel = self.poly.inputs[0].size
def__call__(self, weights, node_feats, sph, senders, receivers, num_nodes):
# weights: (num_edges, weight_numel)
w = rearrange(weights, "e (s m) -> e s m", s=self.poly.inputs[0].num_segments)
# node_feats: dict[Irrep, Array] with (nodes, mul, ir.dim)
x1 = jax.tree.map(lambda v: rearrange(v, "n m i -> n i m"), node_feats)
# sph: dict[Irrep, Array] with (edges, 1, ir.dim) or (edges, ir.dim)
x2 = jax.tree.map(lambda v: rearrange(v, "e 1 i -> e i"), sph)
out_template = {
ir: jax.ShapeDtypeStruct(
(num_nodes, desc.num_segments) + desc.segment_shape, w.dtype
)
for (_, ir), desc inzip(self.irreps_out, self.poly.outputs)
}
y = cuex.ir_dict.segmented_polynomial_uniform_1d(
self.poly, [w, x1, x2], out_template,
input_indices=[None, senders, None],
output_indices=receivers,
name="tensor_product",
)
# Rearrange output back to (nodes, mul, ir.dim) for downstream layersreturn {
ir: rearrange(v, "n (i s) m -> n (s m) i", i=ir.dim)
for ir, v in y.items()
}
RepArray
Representation-aware JAX array:
rep = cue.IrrepsAndLayout(cue.Irreps("SO3", "4x0 + 2x1"), cue.ir_mul)
x = cuex.RepArray(rep, jnp.ones((batch, rep.dim)))
x = cuex.randn(jax.random.key(0), rep, (batch,), jnp.float32)
x.array # raw jax.Array
x.reps # {axis: Rep}
x.irreps # Irreps (if last axis is IrrepsAndLayout)
