Research map / Probabilistic and scientific modeling / Dynamical systems
JAX MD: A Framework for Differentiable Physics
Read the paperCheck your idea against 108,967 papers
AI-extracted summary. Written by AI from the paper's title only (no summary yet). Numbers and claims can be wrong or missing context, so check the paper before you cite or build on it.
Measured effects reported in the paper
| What was changed | Where | Measured | Effect |
|---|---|---|---|
| run JAX MD instead of HOOMD-Blue on small Lennard-Jones system (GPU) | JAX MD · LJ N=216 K80 GPU | ms/step: 0.07 -> 0.09 | No clear effect |
| run JAX MD instead of LAMMPS on small Lennard-Jones system (GPU) | JAX MD · LJ N=216 K80 GPU | ms/step: 0.20 -> 0.09 | Helped |
| run JAX MD instead of LAMMPS on large Lennard-Jones system (GPU) | JAX MD · LJ N=64k K80 GPU | ms/step: 4.0 -> 10.5 | Hurt |
| run JAX MD instead of HOOMD-Blue on large Lennard-Jones system (GPU) | JAX MD · LJ N=64k K80 GPU | ms/step: 4.5 -> 10.5 | Hurt |
| run JAX MD instead of LAMMPS on large system (CPU) | JAX MD · LJ N=10k CPU | ms/step: 120 -> 700 | Hurt |
| run JAX MD instead of LAMMPS on small system (CPU) | JAX MD · LJ N=216 CPU | ms/step: 0.1 -> 1.9 | Hurt |
| move JAX MD from GPU to TPU on small system | JAX MD · LJ N=216 | ms/step: 0.09 -> 0.02 | Helped |
| move JAX MD from TPU small to large system | JAX MD on TPU · LJ N=10k TPU | ms/step: 0.02 -> 360 | Hurt |
| train graph network potential on DFT energies and forces with rotation/flip augmentation | GNN energy model · Silicon DFT test set | energy RMSE (meV/atom): 2.8 | Helped |
Closest papers by meaning
- Long-Short-Range Message-Passing: A Physics-Informed Framework to Capture Non-Local Interaction for Scalable Molecular Dynamics Simulation (ICLR 2024 · similarity 0.85)
- $\bf{\Phi}_\textrm{Flow}$: Differentiable Simulations for PyTorch, TensorFlow and Jax (ICML 2024 · similarity 0.85)
- Mosaic: A Benchmark Suite for Differentiable Physics Solvers (NeurIPS 2026 · similarity 0.84)
- MaNGO — Adaptable Graph Network Simulators via Meta-Learning (NeurIPS 2025 · similarity 0.84)
- Neural Network Reparametrization for Accelerated Optimization in Molecular Simulations (NeurIPS 2024 · similarity 0.84)
- Towards Fast, Specialized Machine Learning Force Fields: Distilling Foundation Models via Energy Hessians (ICLR 2025 · similarity 0.84)
Writing a paper in this area? scime draft is a free LaTeX editor with AI included, real-time co-authors and a check of your references and venue rules before you submit. It also imports OpenAI Prism projects.