neural_lam.utils.networks

Constructors for neural-network building blocks.

Functions

make_gnn_seq(→ torch.nn.Module)	Build a sequential stack of GNN layers that propagates both node and
make_mlp(→ torch.nn.Sequential)	Construct a multilayer perceptron from a blueprint of layer widths.

Module Contents

neural_lam.utils.networks.make_gnn_seq(edge_index: torch.Tensor, num_gnn_layers: int, hidden_layers: int, hidden_dim: int, gnn_type: str = 'InteractionNet') → torch.nn.Module

Build a sequential stack of GNN layers that propagates both node and edge representations.

All layer types share the (send, rec, edge) -> (rec, edge) interface, so the stack can be applied as a single module.

Parameters:

• edge_index (torch.Tensor) – Shape (2, M). Edge index of the edges that the GNN layers operate on.

• num_gnn_layers (int) – Number of stacked GNN layers; must be at least 1. Callers that want a no-op stage (e.g. zero intra-level layers) should skip building and applying the stack rather than calling this with 0.

• hidden_layers (int) – Number of hidden layers in the MLPs of each GNN layer.

• hidden_dim (int) – Dimensionality of node and edge representations.

• gnn_type (str) – GNN layer type, any key in gnn_layers.GNN_TYPES.

Returns:

Sequential module mapping (mesh_rep, edge_rep) to updated (mesh_rep, edge_rep).

Return type:

pyg.nn.Sequential

Raises:

ValueError – If num_gnn_layers is less than 1.

neural_lam.utils.networks.make_mlp(blueprint: list[int], layer_norm: bool = True) → torch.nn.Sequential

Construct a multilayer perceptron from a blueprint of layer widths.

Parameters:

• blueprint (list[int]) – Sequence of layer dimensions where blueprint[0] is the input size, blueprint[-1] is the output size, the intermediate entries specify the hidden layer widths, and len(blueprint) - 2 is the number of hidden layers.

• layer_norm (bool, optional) – If True, append a LayerNorm to the output as in GraphCast.

Returns:

Sequential module implementing the specified MLP.

Return type:

torch.nn.Sequential