ProgramMaster Logo
Conference Tools for Materials Science & Technology 2020
Register as a New User
Submit An Abstract
Propose A Symposium
Presenter/Author Tools
Organizer/Editor Tools
About this Abstract
Meeting Materials Science & Technology 2020
Symposium Ceramics and Glasses Simulations and Machine Learning
Presentation Title JAX, M.D.: End-to-End Differentiable, Hardware Accelerated, Molecular Dynamics in Pure Python
Author(s) Samuel Schoenholz, Ekin Dogus Cubuk
On-Site Speaker (Planned) Samuel Schoenholz
Abstract Scope Molecular Dynamics (MD) software is used across a vast range of subjects from physics and materials science to biochemistry and drug discovery. Most MD software involves significant use of handwritten derivatives and code reuse across C++, FORTRAN, and CUDA. In this work we bring the substantial advances in software that have taken place in machine learning to MD with JAX, M.D. (JAX MD). JAX MD is an end-to-end differentiable MD package written entirely in Python that can be just-in-time compiled to CPU, GPU, or TPU. JAX MD allows researchers to iterate extremely quickly and lets researchers easily incorporate machine learning models into their workflows. In addition to making workloads easier, JAX MD allows researchers to take derivatives through whole-simulations to design Physical systems with desirable properties. We discuss the architecture of JAX MD through several vignettes with an eye towards glass physics. Code available at
Proceedings Inclusion? Undecided


Ab-initio and Reactive MD Simulations of Polymer Pyrolysis and Formation of Silicon-based Ceramics
Application of Natural Language Processing to Zeolites and Cementitious Materials
Beyond the Average: Fluctuations in Glass-forming Systems
Data, Materials and Disorder
De Novo Discovery of Nanoporous Structures with Tailored Sorption Isotherm by Machine Learning
Defect Formation and Self-diffusion in Alumina: Computational Approaches
Introductory Comments: Ceramics and Glasses Simulations and Machine Learning
JAX, M.D.: End-to-End Differentiable, Hardware Accelerated, Molecular Dynamics in Pure Python
The Energy Landscape Governs Brittle-to-Ductile Transitions in Glasses
The Role of Pore Pattern on The Ductility Enhancement of Crystalline Silicon Nitride Nanoporous Membranes
Theoretical Calculation of Formation Energies and Site Preference of Substitutional Divalent Cations in Carbonated Apatite
Verification of Mn Local Structure in Manganese Lithium Borate-based Glass by Computer Simulations and X-ray Absorption Spectroscopy

Questions about ProgramMaster? Contact