DeepMind on JAX

DT #59 — February 14, 2021

AI lab tooling long read #1: DeepMind published a blog post about using JAX to accelerate their research. JAX is a modern take on the NumPy API that “includes an extensible system of composable transformation that help support machine learning research” by taking care of differentiation, vectorization (like abstracting batching away from the researcher), and JIT-compilation (for GPUs and TPUs). The Python library now underpins many of DeepMind’s recent publications, and they’ve also open-sourced several components of their internal ecosystem on top of JAX: Haiku, Optax, RLax, Chex, and Jraph (“it’s pronounced gif ”).

ML Research

This section of Dynamically Typed covers recent models, datasets, and tools for machine learning research.

Join 325+ others and subscribe to get DT in your inbox every second Sunday — 76 issues and counting!

Or check out recent DT issues first:

DT #76: Dynamically Typed Hiatus

DT #75: OpenAI's book summaries for the alignment problem, Translatotron 2, and AI-generated movie posters

DT #74: Apple's privacy-focused facial recognition, DeepMind's multimodal Perceiver IO, and sea ice forecasting with IceNet