Matthew Johnson

Timeline

1. 2010–2014

   Bayesian time series at MIT

   After undergraduate study at UC Berkeley, he earns a PhD in EECS at MIT, where he works with Alan Willsky on Bayesian time series models and scalable inference. With Willsky, he introduces the explicit-duration hierarchical Dirichlet process hidden semi-Markov model at UAI in 2010, and in 2014 he presents stochastic variational inference algorithms for hidden Markov and semi-Markov models at ICML.

2. 2015

   Mouse behavior at Harvard

   He works as a postdoc with Ryan Adams in the Harvard Intelligent Probabilistic Systems Group and with Bob Datta in the Datta Lab at Harvard Medical School. In 2015, he co-authors a Neuron paper with Alexander Wiltschko and colleagues that combines 3D depth imaging and machine learning to describe mouse behavior as a series of reused, stereotyped modules.

3. March 2016

   Graphical models with neural networks

   In March 2016, he and David Duvenaud, Alexander Wiltschko, Sandeep R. Datta and Ryan Adams propose a framework that composes probabilistic graphical models with deep learning methods. Their model family adds neural network observation models to graphical latent structure, and its inference extends variational autoencoders.

4. 2018

   The JAX tracing compiler

   He co-authors a 2018 SysML paper with Roy Frostig and Chris Leary that describes JAX, a tracing JIT compiler that turns pure Python and NumPy programs into accelerator code through XLA. JAX is fully compatible with Autograd, a precursor library that he helped found and co-wrote with Dougal Maclaurin, David Duvenaud and Jamie Townsend.

5. December 2018

   Open-sourcing JAX

   JAX has its initial open-source release in December 2018. He is one of its authors, with James Bradbury, Roy Frostig, Peter Hawkins, Chris Leary, Dougal Maclaurin and others.

6. 2020

   JAX tech lead at Google Brain

   In November 2020, he presents JAX at the REBASE track of SPLASH as a research scientist at Google Brain. He also serves as tech lead for JAX, whose composable transformations include automatic differentiation, automatic batching, compilation through XLA and parallelization across accelerators.

7. April 2022

   Decomposing reverse-mode differentiation

   In April 2022, he co-authors You Only Linearize Once with Alexey Radul, Adam Paszke, Roy Frostig and Dougal Maclaurin. Following the automatic differentiation systems in JAX and Dex, the paper splits reverse-mode differentiation into forward-mode differentiation, unzipping of linear and non-linear parts, and transposition.

8. 2025

   Open-source tools at Google DeepMind

   By 2025, he is a researcher at Google DeepMind, where he develops open-source tools and frameworks for machine learning and is a key contributor to JAX. In July 2025, he gives the talk “JAX and OSS at Google / DeepMind” at the CODEML workshop at ICML 2025 in Vancouver.