JaxPM

JAX-powered Cosmological Particle-Mesh N-body Solver

The goal of this project is to build the next generation of differentiable PM solvers, using JAX as the backend.