based on our open-source HhPF simulator HyMD. In ∂-HyMD we use JAX to trace the update steps and get the gradient of the loss function as in Equation , while implementing the HhPF operations with fast Fourier transforms, the MD integrator, barostat and thermostat using JAX NumPy API, and taking advantage of JIT compilations whenever possible.