Building the Energy-Conserving Net in JAX/Flax NNX

A runnable companion to the energy-conservation post: the HNN pendulum field as the symplectic gradient of one learned scalar, the plain field model it beats on drift, and the leapfrog classifier whose residual block is a kick-drift-kick step of a learned potential, all as Flax NNX modules with lax.scan doing depth. Every figure is rendered from the real Kaggle run.