waxmorph.jax.gnn

Equinox encode-process-decode GNS (Sanchez-Gonzalez et al., ICML 2020).

Static JIT shapes use padded directed contact edges; num_edges masks padding before message aggregation. This is the behavioral counterpart of waxmorph.torch.gnn.

Classes

GNS(node_feature_dim, edge_feature_dim[, ...])	Encode graph features, apply residual message passing, and decode state deltas.
GraphNetworkBlock([node_latent_dim, ...])	Residual edge and node updates over directed COO messages.