xnn.gnn.fast.symmetric_contraction.SymmetricContractionKernel

class xnn.gnn.fast.symmetric_contraction.SymmetricContractionKernel(sc, device, dtype)

Bases: object

cuEquivariance evaluation of one SymmetricContraction on one device and dtype.

Parameters:

• sc (xnn.gnn.models.mace.SymmetricContraction) – The reference module (its U buffers define the basis of its weights).

• device (torch.device) – The CUDA device the kernel runs on.

• dtype (torch.dtype) – float32 or float64.