waxmorph.jax.gnn.GNS

waxmorph.jax.gnn.GNS#

class waxmorph.jax.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, *, key)[source]#

Bases: Module

Encode graph features, apply residual message passing, and decode state deltas.

Default heads dX, dP, and dc are raw per-node position, polarity, and concentration deltas; training scales them by dt_gns. key is split across all submodules. checkpoint_processor rematerializes processor blocks and decoder heads, trading compute for fewer saved activations.

Examples

>>> model = GNS(1, 2, hidden_dim=4, num_mp_steps=1, output_dims={"dX": 3}, key=jax.random.PRNGKey(0))
>>> out = model(jnp.ones((2, 1)), jnp.array([[0, 1], [1, 0]]), jnp.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)

  • key (jax.Array)

__call__(node_features, edge_index, edge_features, num_edges=None)[source]#

Decode per-node heads from a directed COO graph.

num_edges masks padded rows. Each value has shape [N, output_dims[head]].

Parameters:
Return type:

dict[str, Array]

save(path)[source]#

Write Equinox leaves to path and constructor config to path.json.

Parameters:

path (str | Path)

Return type:

None

classmethod load(path)[source]#

Reconstruct from the path.json config and leaves stored at path.

Parameters:

path (str | Path)

Return type:

GNS