xnn.gnn.models.nequip.ConvNetLayer#

class xnn.gnn.models.nequip.ConvNetLayer(irreps_in, feature_irreps_hidden, irreps_node_attr, irreps_edge_attr, n_radial, resnet=False, nonlinearity_scalars={'e': 'silu', 'o': 'tanh'}, nonlinearity_gates={'e': 'silu', 'o': 'tanh'}, **conv_kwargs)[source]#

Bases: Module

One NequIP layer: InteractionBlock + gated nonlinearity (+ resnet).

Mirrors nequip.nn.ConvNetLayer: the desired hidden irreps are pruned to those reachable by a tensor-product path from the current features and the edge attributes, split into scalars (activated directly) and gated irreps (multiplied by activated scalar gates), and the convolution outputs the gate’s input irreps. The residual (“resnet”) update is applied only when the layer preserves the feature irreps.

Parameters:

irreps_in (e3nn.o3.Irreps) – Irreps of the incoming node features.
feature_irreps_hidden (e3nn.o3.Irreps) – The desired hidden irreps (see nequip_hidden_irreps()).
irreps_node_attr (e3nn.o3.Irreps) – Irreps of the one-hot species node attributes.
irreps_edge_attr (e3nn.o3.Irreps) – Irreps of the spherical-harmonic edge attributes.
n_radial (int) – Width of the invariant radial edge embedding.
resnet (bool, optional) – Residual update when the irreps allow it, by default False.
nonlinearity_scalars (dict, optional) – Per-parity activation names (upstream defaults {"e": "silu", "o": "tanh"}).
nonlinearity_gates (dict, optional) – Per-parity activation names (upstream defaults {"e": "silu", "o": "tanh"}).
**conv_kwargs – Forwarded to InteractionBlock (invariant_layers, invariant_neurons, avg_num_neighbors, use_sc).

Variables:

conv (InteractionBlock) – The convolution (named conv to match upstream state dicts).
equivariant_nonlin (_Gate) – The gated nonlinearity.
irreps_out (e3nn.o3.Irreps) – Irreps of the output node features.

forward(x, node_attrs, edge_index, edge_sh, edge_radial)[source]#

Apply one message-passing layer.

Parameters:

x (torch.Tensor) – Node features of shape (N, irreps_in.dim).
node_attrs (torch.Tensor) – One-hot species node attributes.
edge_index (torch.Tensor) – Edge index of shape (2, E).
edge_sh (torch.Tensor) – Spherical-harmonic edge attributes.
edge_radial (torch.Tensor) – Invariant radial edge embedding.

Returns:

Updated node features of shape (N, irreps_out.dim).

Return type:

torch.Tensor