Tuesday 20th October 14.00-15.00 B2.02 (Sci Conc)
Dear Cohorts 4, 5, 6 and 7 -
Please join us for this informal seminar, followed by coffee - sign up below so that we have an idea of numbers please ![]()
https://warwick.ac.uk/fac/sci/hetsys/newsandevents/seminar/
Speaker: Abhijeet Gangan
Title: Molecular dynamics in differentiable languages: lessons from JAX-MD and TorchSim
Abstract: JAX and PyTorch are popular languages for ML applications, but they can also be used to build simulation engines for molecular dynamics. In this presentation, I will discuss my experience developing JAX-MD and TorchSim to show what this approach provides: automatic differentiation, batching, just-in-time compilation and native support for machine-learned potentials.
I will also discuss the limitations, such as static-shape array operations being a poor fit for dynamic neighbour lists and variable system sizes, which can lead to recompilation, memory overhead and irregular memory access. Finally, I will outline the path ahead for performance, covering topics like custom kernels and coupling JAX/PyTorch models to production MD codes such as LAMMPS and OpenMM.
Bio: I currently work as a Member of Technical Staff at Periodic Labs. Before that, I completed my PhD at UCLA, during which I also worked at Radical AI and Lawrence Livermore National Laboratory. My work is focused on developing atomistic simulation codes and machine-learned interatomic potentials.