EI Seminar - Matthew Johnson - JAX: accelerated ML research via composable function transformations 

Speaker: Matthew Johnson
Title: JAX: accelerated machine learning research via composable function transformations in Python
Abstract: This talk is about JAX, a system for high-performance machine learning research and numerical computing. It offers the familiarity of Python+NumPy together with hardware acceleration. JAX combines these features with user-wielded function transformations, including automatic differentiation, automatic vectorized batching, end-to-end compilation (via XLA), parallelizing over multiple accelerators, and more. Composing these transformations is the key to JAX's power and simplicity.
Bio: Matt Johnson is a research scientist at Google Brain interested in software systems powering machine learning research. He's the tech lead for JAX. When moonlighting as a machine learning researcher, he works on making neural ODEs faster to solve, automatically exploiting conjugacy in probabilistic programs, and composing graphical models with neural networks. Matt was a postdoc with Ryan Adams at the Harvard Intelligent Probabilistic Systems Group and Bob Datta in the Datta Lab at the Harvard Medical School. His Ph.D. is from MIT in EECS, where he worked with Alan Willsky at LIDS on Bayesian time series models and scalable inference.