Apax: Atomistic learned Potentials in JAX!

Apax is a high-performance, extendable package for training of and inference with atomistic neural networks., most prominently the Gaussian Moment Neural Network. It is based on JAX and uses JaxMD as a molecular dynamics engine.