waxmorph.torch.gnn.GNS

waxmorph.torch.gnn.GNS#

class waxmorph.torch.gnn.GNS(node_feature_dim, edge_feature_dim, node_latent_dim=128, edge_latent_dim=128, hidden_dim=128, num_mp_steps=10, num_mlp_layers=2, output_dims=None, activation='relu', layer_norm=True, checkpoint_processor=False)[source]#

Bases: Module

Encode-process-decode graph network.

Independent residual message-passing blocks consume encoded node and edge features. Named decoders map final node latents to per-node updates; defaults are dX: 3, dP: 3, and dc: 2. num_mlp_layers counts linear layers in every MLP. layer_norm applies to encoders and processor MLPs; decoders project directly to outputs. checkpoint_processor recomputes processor and decoder activations during backward to reduce saved activation memory.

Examples

>>> model = GNS(1, 2, hidden_dim=4, num_mp_steps=1, output_dims={"dX": 3})
>>> out = model(torch.ones(2, 1), torch.tensor([[0, 1], [1, 0]]), torch.ones(2, 2))
>>> print(sorted(out), tuple(out["dX"].shape))
['dX'] (2, 3)
Parameters:

  • node_feature_dim (int)

  • edge_feature_dim (int)

  • node_latent_dim (int)

  • edge_latent_dim (int)

  • hidden_dim (int)

  • num_mp_steps (int)

  • num_mlp_layers (int)

  • output_dims (dict[str, int] | None)

  • activation (str)

  • layer_norm (bool)

  • checkpoint_processor (bool)

save(path)[source]#

Serialize {"config": ..., "state_dict": ...} with torch.save().

Parameters:

path (str | Path)

Return type:

None

classmethod load(path, **kwargs)[source]#

Load the config/state_dict schema written by save().

Use checkpoints from trusted sources; weights_only=False can execute serialized Python code. kwargs pass to torch.load().

Parameters:

path (str | Path)

Return type:

GNS

forward(node_features, edge_index, edge_features)[source]#

Decode per-node updates from raw graph features.

edge_index is directed COO [2, E] with senders in row 0. Each configured output head maps to [N, output_dim].

Parameters:
Return type:

dict[str, Tensor]