xnn.gnn.fast.conv#
The message step of an equivariant convolution, with a fused fast path.
Equivariant message passing (MACE, NequIP and their relatives) spends most of
its time in one step: gather the sender features onto the edges, take the
uvu tensor product with the edge spherical harmonics weighted per edge by a
radial network, and sum the messages onto the receivers. The reference does the
three as separate operations and keeps the per-edge intermediates for the
backward pass; the fast path runs them as one cuEquivariance kernel that reads
the node features and writes the node sums directly.
Classes
|
|