waxmorph.jax

Modules

waxmorph.jax.gnn	Equinox encode-process-decode GNS (Sanchez-Gonzalez et al., ICML 2020).
waxmorph.jax.graph	JAX graph construction with detached topology and differentiable features.
waxmorph.jax.losses	JAX losses for ordered arrays and unordered point clouds.
waxmorph.jax.mlp	Equinox MLP used by the JAX graph-network simulator.
waxmorph.jax.warp_autograd	JAX/Warp custom-VJP bridge for frozen-pair physics on CUDA.

Explicit JAX/Equinox backend; top-level waxmorph imports use PyTorch.

JAX uses padded static graph buffers, Optax, explicit PRNG keys, and custom VJPs on CUDA for Warp physics. Its loss factory provides OTT Sinkhorn divergence; PyTorch provides GeomLoss families.