DeepMind: Using JAX to accelerate our research

Recently, we've found that an increasing number of projects are well served by JAX, a machine learning framework developed by Google Research teams. JAX resonates well with our engineering philosophy and has been widely adopted by our research community over the last year. Here we share our experience of working with JAX, outline why we find it useful for our AI research, and give an overview of the ecosystem we are building to support researchers everywhere.

When DeepMind adopts a tool, it's worth paying attention.


Want to receive more content like this in your inbox?