xnn.gnn.models.mace.RealAgnosticInteractionBlock#

class xnn.gnn.models.mace.RealAgnosticInteractionBlock(node_attrs_irreps, node_feats_irreps, edge_attrs_irreps, edge_feats_irreps, target_irreps, hidden_irreps, avg_num_neighbors, radial_MLP)[source]#

Bases: _InteractionBase

Non-residual interaction: the skip connection is applied to the message.

A _InteractionBase whose skip_tp mixes the aggregated, normalised message with the node element attributes (so the self-connection acts on the message rather than the input features). Returns None in place of a separate self-connection.

Parameters:

avg_num_neighbors (float)

forward(node_attrs, node_feats, edge_attrs, edge_feats, edge_index)[source]#

Compute one message-passing update (non-residual variant).

Parameters:

  • node_attrs (torch.Tensor) – Per-node one-hot element attributes.

  • node_feats (torch.Tensor) – Incoming node features, shape (num_nodes, node_feats_irreps.dim).

  • edge_attrs (torch.Tensor) – Edge spherical-harmonic attributes.

  • edge_feats (torch.Tensor) – Scalar radial edge features feeding the radial MLP.

  • edge_index (torch.Tensor) – Edge index of shape (2, num_edges) ([senders, receivers]).

Returns:

The reshaped message features and None (no separate self-connection for the non-residual block).

Return type:

tuple of (torch.Tensor, None)