JAX is a system for high-performance machine learning research and numerical computing. It offers the familiarity of Python+NumPy together with a collection of user-wielded function transformations, including automatic differentiation, vectorized batching, end-to-end compilation (via XLA), parallelization over multiple accelerators, and more. Composing these transformations is the key to JAX’s power and simplicity.

This talk presents an overview of the project today and highlights some of our discoveries so far. These include useful autodiff primitives as well as lessons from embedding strictly-typed pure functional languages within the unruly realm of research programming in Python.