xnn.gnn.models.mace.RealAgnosticResidualInteractionBlock#

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

Bases: _InteractionBase

Residual interaction: self-connection computed from the input features.

A _InteractionBase whose skip_tp builds a self-connection from the input node features and element attributes. That self-connection is returned alongside the message so the downstream product basis can add it as a residual.

Parameters:

avg_num_neighbors (float)

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

Compute one message-passing update (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 the self-connection sc computed from the input features (to be added as a residual).

Return type:

tuple of (torch.Tensor, torch.Tensor)