AI People Wiki

Builders & open source

Matthew Johnson

Research scientist (GitHub bio "research scientist @ Google Brain"), Google

  • Research scientist (GitHub bio "research scientist @ Google Brain"), Google52

Matthew Johnson, listed as Matthew James Johnson, is one of the authors named in the JAX repository's citation entry. With Roy Frostig and Chris Leary, Johnson co-authored the SysML 2018 paper "Compiling machine learning programs via high-level tracing," which the JAX README identifies as describing a nascent version of JAX; the paper lists Johnson with Google Brain. Johnson also wrote JAX design documents published in its documentation, including the Omnistaging upgrade guide (2020), and co-authored the January 2023 design doc that proposed shard_map. Johnson's GitHub profile lists Google as their company.12345

At a glance

Current role
Research scientist (GitHub bio "research scientist @ Google Brain")
Company
Google
Sector
Company
Work listed
1
Sources cited
9
Profile reviewed

Research

Matthew Johnson's research has combined probabilistic modeling with deep learning. In 2016 Johnson was first author, with David Duvenaud, Alexander Wiltschko, Sandeep Datta and Ryan Adams, of "Composing graphical models with neural networks for structured representations and fast inference", which joins probabilistic graphical models with neural network observation models.6

Johnson has also worked on automatic differentiation. With Alexey Radul, Adam Paszke, Roy Frostig and Dougal Maclaurin, Johnson co-authored "You Only Linearize Once: Tangents Transpose to Gradients" (2022), which follows up on the differentiation systems in the JAX and Dex projects by formalizing reverse mode as forward mode followed by unzipping and transposition.7

Open-source work

Johnson is one of the authors of Autograd, the Python library for automatically differentiating NumPy code, whose README credits Dougal Maclaurin, David Duvenaud, Matt Johnson and Jamie Townsend as its writers. Johnson's GitHub profile describes Johnson as a research scientist at Google Brain.85

Johnson then co-authored, with Roy Frostig and Chris Leary, the SysML 2018 paper "Compiling machine learning programs via high-level tracing", which the JAX README identifies as describing a nascent version of JAX, and is named in the repository's citation entry. In JAX's design documents Johnson wrote the September 2020 guide to the omnistaging upgrade and co-authored the January 2023 proposal for shard_map.21934

Work and contributions

  • JAX Open-source code

    Named in the "Citing JAX" author list; co-author of the SysML 2018 paper on an early version of JAX; author of the Omnistaging guide and a co-author of the shard_map design doc in the JAX documentation. The repository is released under the Apache License 2.0.12349

Timeline

  1. First author of a paper on composing graphical models with neural networks for structured representations.6
  2. Co-authored the SysML 2018 paper describing an early version of JAX with Roy Frostig and Chris Leary.2
  3. Wrote the JAX design document on the omnistaging upgrade.3
  4. Co-authored "You Only Linearize Once: Tangents Transpose to Gradients".7
  5. Co-authored the JAX design document proposing shard_map.4

Where they work

Suggest a correction or removal

Profiles cover public professional work only, and each statement cites a source. If something is out of date, missing or wrong, tell us and link a source. People can also ask for their profile to be corrected or removed.

References

9 sources · profile reviewed Oct 21, 2026 · data updated Oct 21, 2026

  1. jax-ml/jax README ("Citing JAX" author list)jax-ml (GitHub repository) · Repository · accessed Oct 21, 2026
  2. Compiling machine learning programs via high-level tracing (SysML 2018)MLSys · Paper · published 2018 · accessed Oct 21, 2026
  3. Omnistaging (JAX design note, September 2020)JAX documentation · Documentation · published Sep 25, 2020 · accessed Oct 21, 2026
  4. shmap (shard_map) for simple per-device code (JAX design doc, January 2023)JAX documentation · Documentation · published Jan 2023 · accessed Oct 21, 2026
  5. mattjj (Matthew Johnson) on GitHubGitHub · Official page · accessed Oct 21, 2026
  6. Composing graphical models with neural networks for structured representations and fast inference (arXiv 1603.06277)arXiv · Paper · published Mar 20, 2016 · accessed Oct 21, 2026
  7. You Only Linearize Once: Tangents Transpose to Gradients (arXiv 2204.10923)arXiv · Paper · published Apr 22, 2022 · accessed Oct 21, 2026
  8. HIPS/autograd READMEHIPS (GitHub repository) · Repository · accessed Oct 21, 2026
  9. jax-ml/jax LICENSE (Apache License 2.0)jax-ml (GitHub repository) · License · accessed Oct 21, 2026