xnn.gnn.models.blocks.tp_out_irreps_with_instructions#

xnn.gnn.models.blocks.tp_out_irreps_with_instructions(irreps1, irreps2, target_irreps, sort_instructions=True)[source]#

(uvu) tensor-product output irreps + instructions, keeping only target paths.

Enumerates the uvu tensor-product paths between irreps1 and irreps2, keeping only those whose output irrep lies in target_irreps, then sorts the output irreps and remaps the instruction indices accordingly. Used to build the convolution e3nn.o3.TensorProduct in the MACE and NequIP interaction blocks.

Parameters:

irreps1 (e3nn.o3.Irreps) – First operand irreps (node features).
irreps2 (e3nn.o3.Irreps) – Second operand irreps (edge spherical-harmonic attributes).
target_irreps (e3nn.o3.Irreps) – Only output irreps present in this set are retained.
sort_instructions (bool, optional) – Also sort the instructions by output index (the mace-torch convention; the original nequip keeps enumeration order). The choice fixes the tensor-product weight layout, so it must match the upstream code weights are transplanted from. Default is True.

Returns:

The sorted output irreps and the corresponding list of tensor-product instructions (i, j, k, "uvu", True) referencing the sorted indices.

Return type:

tuple of (e3nn.o3.Irreps, list)