xnn.gnn.models.mace.SymmetricContraction#

class xnn.gnn.models.mace.SymmetricContraction(irreps_in, irreps_out, correlation, num_elements)[source]#

Bases: Module

Per-element symmetric contraction over all output irreps (the MACE product basis).

The learned, higher-body-order product basis of MACE (Batatia et al. 2022, Eq. 10-11). It holds one _Contraction per output irrep and concatenates their results, reproducing mace-torch to ~1e-16 given the same weights.

Parameters:

  • irreps_in (e3nn.o3.Irreps) – Input node-feature irreps fed to every per-irrep contraction.

  • irreps_out (e3nn.o3.Irreps) – Target output irreps; one _Contraction is created per entry.

  • correlation (int) – Maximum correlation order (body order minus one).

  • num_elements (int) – Number of chemical elements (weights are element-dependent).

Variables:

contractions (torch.nn.ModuleList) – One _Contraction per output irrep.

forward(x, y)[source]#

Run every per-irrep contraction and concatenate the results.

Parameters:

  • x (torch.Tensor) – Node features of shape (B, num_features, coupling_dim).

  • y (torch.Tensor) – Per-node one-hot element attributes of shape (B, num_elements).

Returns:

The concatenated contracted features spanning all output irreps.

Return type:

torch.Tensor