xnn.gnn.fast#

Fast paths of the equivariant GNN blocks (fused cuEquivariance kernels).

ConvTensorProduct (the message step of MACE and NequIP) and the symmetric contraction of MACE have a fused GPU implementation next to their reference one; FastPathModule explains the contract and use_fast the choice. AUTO_POLICY holds the measured size above which use_fast="auto" takes the kernels.

Module Attributes

AUTO_POLICY

Edge counts above which use_fast="auto" uses the kernels.

Modules

`conv`	The message step of an equivariant convolution, with a fused fast path.
`symmetric_contraction`	The fast path of the MACE symmetric contraction (the product basis).