xnn.gnn.fast.conv.ConvTensorProduct#

class xnn.gnn.fast.conv.ConvTensorProduct(*args, **kwargs)[source]#

Bases: TensorProduct, FastPathModule

uvu tensor product with per-edge weights and a fused message pass.

An e3nn.o3.TensorProduct (same arguments, parameters and state_dict) with one more method, conv(), which evaluates the whole message step. Called as a module it is the plain e3nn tensor product.

The fast path applies when every instruction is uvu, the second input has multiplicity one throughout (spherical harmonics) and the weights are per edge (shared_weights=False), which is the convolution of MACE and NequIP. It is exact up to the summation order.

fast_supported(device, dtype)[source]#

Whether the fast path can run for inputs on device in dtype.

Parameters:
Return type:

bool

conv(node_feats, edge_attrs, edge_weights, edge_index, num_nodes, message_divisor=None)[source]#

Gather, tensor product and sum onto the receivers in one call.

Parameters:

  • node_feats (torch.Tensor) – Node features, (N, irreps_in1.dim).

  • edge_attrs (torch.Tensor) – Edge attributes (spherical harmonics), (E, irreps_in2.dim).

  • edge_weights (torch.Tensor) – Per-edge path weights, (E, weight_numel).

  • edge_index (torch.Tensor) – (2, E): row 0 the senders, row 1 the receivers.

  • num_nodes (int) – N.

  • message_divisor (float, optional) – Divide every message by this before summing (NequIP’s sqrt(avg_num_neighbors)); the fast path folds it into the weights.

Returns:

The summed messages, (N, irreps_out.dim).

Return type:

torch.Tensor