xnn.gnn.models.mace.RealAgnosticDensityResidualInteractionBlock#

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

Bases: RealAgnosticResidualInteractionBlock

Residual interaction with learned density normalization.

Identical to RealAgnosticResidualInteractionBlock except for the 1 + rho_i message normalization of RealAgnosticDensityInteractionBlock.

Parameters:

avg_num_neighbors (float)

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

Compute one density-normalized 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)