Informal Seminar: Molecular Dynamics in Differentiable Languages

Sarah Jarratt — Wed, 14 Oct 2026 14:27:02 GMT

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.

Heterogeneous Systems » HetSys Cohort Newsfeed » Informal Seminar: Molecular Dynamics in Differentiable Languages

Informal Seminar: Molecular Dynamics in Differentiable Languages