JAX, M.D. A framework for differentiable physics*

DOI10.1088/1742-5468/ac3ae9OpenAlexW4205254860MaRDI QIDQ5020064

Ekin D. Cubuk, Samuel S. Schoenholz

Publication date: 3 January 2022

Published in: Journal of Statistical Mechanics: Theory and Experiment (Search for Journal in Brave)

Full work available at URL: https://arxiv.org/abs/1912.04232

zbMATH Keywords

molecular dynamicsnumerical simulationsmachine learning

Mathematics Subject Classification ID

Statistical mechanics, structure of matter (82-XX)

Related Items

A Model-Constrained Tangent Slope Learning Approach for Dynamical Systems ⋮ JAX MD ⋮ Learning Lagrangian fluid mechanics with E(3)-equivariant graph neural networks

Uses Software

• Torch
• Theano
• NumPy
• TensorFlow
• IPython
• DiffSharp
• OpenMM
• PyTorch
• Autograd
• GitHub
• DeepChem
• JAX
• Omniglot
• Boids
• Haiku
• Taichi
• DiffTaichi
• Zygote
• HOOMD-TF
• SchNet

Cites Work

• Unnamed Item
• Automatic differentiation of explicit Runge-Kutta methods for optimal control
• General purpose molecular dynamics simulations fully implemented on graphics processing units
• Fast parallel algorithms for short-range molecular dynamics
• Human-level concept learning through probabilistic program induction
• On the performance of discrete adjoint CFD codes using automatic differentiation
• Efficient adjoint derivatives: application to the meteorological model meso-nh
• The frontier of simulation-based inference