# JAX, M.D.: A Framework for Differentiable Physics

arXiv: https://arxiv.org/abs/1912.04232 · HTML: https://mlsanity.com/paper/1912.04232/

## Key facts

- Cited by 7 papers in the map (data through 14 Oct 2026).
- First posted on arXiv 9 Dec 2019.

## Built on by (most cited first, 7 of 7)

- [Efficient Differentiable Simulation of Articulated Bodies](https://mlsanity.com/paper/2109.07719/) (2021)
- [Learn2Hop: Learned Optimization on Rough Landscapes](https://mlsanity.com/paper/2107.09661/) (2021)
- [A Comparison of Mesh-Free Differentiable Programming and Data-Driven Strategies for Optimal Control under PDE Constraints](https://mlsanity.com/paper/2310.02286/) (2023)
- [Half-Inverse Gradients for Physical Deep Learning](https://mlsanity.com/paper/2203.10131/) (2022)
- [Stabilizing Backpropagation Through Time to Learn Complex Physics](https://mlsanity.com/paper/2405.02041/) (2024)
- [Unsupervised Discovery of Inertial-Fusion Plasma Physics using Differentiable Kinetic Simulations and a Maximum Entropy Loss Function](https://mlsanity.com/paper/2206.01637/) (2022)
- [Exploring Physical Latent Spaces for High-Resolution Flow Restoration](https://mlsanity.com/paper/2211.11298/) (2022)

## Abstract (arXiv)

We introduce JAX MD, a software package for performing differentiable physics simulations with a focus on molecular dynamics. JAX MD includes a number of physics simulation environments, as well as interaction potentials and neural networks that can be integrated into these environments without writing any additional code. Since the simulations themselves are differentiable functions, entire trajectories can be differentiated to perform meta-optimization. These features are built on primitive operations, such as spatial partitioning, that allow simulations to scale to hundreds-of-thousands of particles on a single GPU. These primitives are flexible enough that they can be used to scale up workloads outside of molecular dynamics. We present several examples that highlight the features of JAX MD including: integration of graph neural networks into traditional simulations, meta-optimization through minimization of particle packings, and a multi-agent flocking simulation. JAX MD is available at www.github.com/google/jax-md.