OverviewHistoryStatsSecurity
npx skills add ...
Documentation
SKILL.md
npx skills add nvidia/cuequivariance --skill cuequivariance-jax
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.
npx skills add nvidia/cuequivariance --skill cuequivariance-jax
cuequivariance_jax (imported as cuex) executes cuequivariance polynomials on GPU via JAX. It provides:
cuex.segmented_polynomial() — JAX primitive with full AD/vmap/JIT supportsegmented_polynomial):
cuex.equivariant_polynomial() + RepArray — the original interface, a single contiguous array with representation metadatacuex.ir_dict module — dict[Irrep, Array] interface, uses IrDictPolynomial descriptors, works naturally with jax.treecuex.nnx module — Flax NNX Module wrappers using dict[Irrep, Array]| Method | Backend | Requirements |
|---|---|---|
"naive" | Pure JAX | Always works, any platform |
"uniform_1d" | CUDA kernel | GPU, all segments uniform shape within each operand, single mode |
"indexed_linear" | CUDA kernel | GPU, linear operations with cuex.Repeats indexing |
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:
Index arrays provide gather (for inputs) and scatter (for outputs). One index per operand (inputs + outputs), None means no indexing:
Fully differentiable — supports jax.grad, jax.jacobian, jax.jvp, jax.vmap:
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.
Use _ir_dict descriptor variants, which return IrDictPolynomial with the polynomial already split by irrep:
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).
The original interface. Wraps segmented_polynomial with RepArray — a single contiguous array with representation metadata:
Equivariant linear layer using dict[Irrep, Array]:
Implementation uses jnp.einsum("uv,...ui->...vi", w, x[ir]) per irrep with 1/sqrt(mul_in) normalization.
Uses spherical_harmonics_ir_dict internally for the dict[Irrep, Array] output:
For species-indexed linear layers (different weights per atom type):
Uses method="indexed_linear" internally with cuex.Repeats.
The uniform_1d CUDA kernel requires:
For equivariant_polynomial() (RepArray interface):
For ir_dict (dict[Irrep, Array] interface), use _ir_dict descriptors directly:
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.
This pattern is used in NequIP, MACE, and similar equivariant GNN models:
Representation-aware JAX array:
| Component | Path |
|---|---|
segmented_polynomial primitive | cuequivariance_jax/segmented_polynomials/segmented_polynomial.py |
uniform_1d backend | cuequivariance_jax/segmented_polynomials/segmented_polynomial_uniform_1d.py |
naive backend | cuequivariance_jax/segmented_polynomials/segmented_polynomial_naive.py |
indexed_linear backend | cuequivariance_jax/segmented_polynomials/segmented_polynomial_indexed_linear.py |
equivariant_polynomial | cuequivariance_jax/equivariant_polynomial.py |
ir_dict module | cuequivariance_jax/ir_dict.py |
nnx module | cuequivariance_jax/nnx.py |
RepArray | cuequivariance_jax/rep_array/rep_array_.py |
Repeats / utilities | cuequivariance_jax/segmented_polynomials/utils.py |
| NequIP example | cuequivariance_jax/examples/nequip_nnx.py |
| MACE example | cuequivariance_jax/examples/mace_nnx.py |
# 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",
)i = jax.random.randint(key, (100, 50), 0, 10) # gather b along axis 0
j1 = jax.random.randint(key, (100, 50), 0, 11) # scatter output axis 0
j2 = jax.random.randint(key, (100, 1), 0, 12) # scatter output axis 1
[out] = cuex.segmented_polynomial(
poly, [a, b, c],
[jax.ShapeDtypeStruct((11, 12, poly.outputs[0].size), jnp.float32)],
indices=[None, np.s_[i, :], None, np.s_[j1, j2]],
method="uniform_1d",
)def loss(w, x, y):
[out] = cuex.segmented_polynomial(
poly, [w, x, y],
[jax.ShapeDtypeStruct((batch, poly.outputs[0].size), jnp.float32)],
method="naive",
)
return jnp.sum(out ** 2)
grad_w = jax.grad(loss, 0)(w, x, y)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 groupfrom 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 in zip(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# Validate dict matches irreps
cuex.ir_dict.assert_mul_ir_dict(irreps, x) # asserts shape (..., mul, ir.dim)
# Convert flat array <-> dict
d = cuex.ir_dict.flat_to_dict(irreps, flat_array) # layout="mul_ir" default
d = cuex.ir_dict.flat_to_dict(irreps, flat_array, layout="ir_mul")
flat = cuex.ir_dict.dict_to_flat(irreps, d)
# Arithmetic
z = cuex.ir_dict.irreps_add(x, y)
z = cuex.ir_dict.irreps_zeros_like(x)
# Create template dict
template = cuex.ir_dict.mul_ir_dict(irreps, jax.ShapeDtypeStruct(shape, dtype))e = cue.descriptors.fully_connected_tensor_product(
4 * cue.Irreps("SO3", "0 + 1"),
cue.Irreps("SO3", "0 + 1"),
4 * cue.Irreps("SO3", "0 + 1"),
)
inputs = [
cuex.randn(jax.random.key(i), rep, (batch,), jnp.float32)
for i, rep in enumerate(e.inputs)
]
# Returns a RepArray with representation metadata
out = cuex.equivariant_polynomial(e, inputs, method="naive")
out.array # the raw jax.Array
out.reps # dict mapping axes to Rep objectsfrom flax import nnx
linear = cuex.nnx.IrrepsLinear(
irreps_in=cue.Irreps(cue.SO3, "4x0 + 2x1").regroup(), # must be regrouped
irreps_out=cue.Irreps(cue.SO3, "3x0 + 5x1").regroup(),
scale=1.0,
dtype=jnp.float32,
rngs=nnx.Rngs(0),
)
# Input/output: dict[Irrep, Array] with shape (batch, mul, ir.dim)
x = {
cue.SO3(0): jnp.ones((batch, 4, 1)),
cue.SO3(1): jnp.ones((batch, 2, 3)),
}
y = linear(x)
# y[cue.SO3(0)].shape == (batch, 3, 1)
# y[cue.SO3(1)].shape == (batch, 5, 3)sh = cuex.nnx.SphericalHarmonics(max_degree=3, eps=0.0)
vectors = jax.random.normal(key, (batch, 3)) # 3D vectors
y = sh(vectors)
# y[cue.O3(0, 1)].shape == (batch, 1, 1) # L=0
# y[cue.O3(1, -1)].shape == (batch, 1, 3) # L=1
# y[cue.O3(2, 1)].shape == (batch, 1, 5) # L=2
# y[cue.O3(3, -1)].shape == (batch, 1, 7) # L=3norm = 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 mulmlp = 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) normalizationindexed_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)e = cue.descriptors.channelwise_tensor_product(...)
e = e.squeeze_modes().flatten_coefficient_modes()
out = cuex.equivariant_polynomial(e, inputs, method="uniform_1d")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]class MessagePassing(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 in zip(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 layers
return {
ir: rearrange(v, "n (i s) m -> n (s m) i", i=ir.dim)
for ir, v in y.items()
}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)