xnn.gnn.fast.symmetric_contraction

The fast path of the MACE symmetric contraction (the product basis).

The contraction is a polynomial of degree correlation in the node features with per-element weights over a coupling basis U. cuEquivariance evaluates the same polynomial space in one fused kernel, with its own coupling basis, so xnn’s weights reach it through a fixed linear change of basis: for each output irrep and order, the kernel’s weights are M @ w with M fitted once, in float64, by matching the two polynomials on random features (exact, as both bases span the same space of equivariant symmetric contractions).

Functions

eligible(sc)	Whether a SymmetricContraction fits the kernel.
fit_basis_map(sc)	The change of basis M (kernel weights = M @ reference weights), in float64.
reference_weights(sc)	The reference weights, per output irrep and then per order (1, 2, ...).

Classes

SymmetricContractionKernel(sc, device, dtype)

cuEquivariance evaluation of one SymmetricContraction on one device and dtype.