waxmorph.jax.train

waxmorph.jax.train#

JAX/Equinox training primitives for shape assembly with a fixed agent count.

`TrainConfig`([n_epochs, t_rollout, ...])	JAX rollout and optimizer settings for a fixed agent count.
`TrainResult`(model, log)	Best post-update model and its run diagnostics.
`train`(model, optimizer, opt_state, loss_fn, ...)	Train a GNS model for shape assembly with a fixed agent count (JAX backend).