Using JAX to accelerate our research by David Budden and Matteo Hessel

deepmind.comsite

Describes the state of JAX and the JAX ecosystem at DeepMind.

from

JAX

added

2026-10-10

likes

0

similar

1. ▲
   Jax in Action manning.com

   site category #jax link ↗

   A hands-on guide to using JAX for deep learning and other mathematically-intensive applications.

2. ▲
   NeurIPS 2020: JAX Ecosystem Meetup youtube.com

   video category #jax link ↗

   JAX, its use at DeepMind, and discussion between engineers, scientists, and JAX core team.

   
3. ▲
   Introduction to JAX by Kevin Murphy colab.research.google.com

   site category #jax link ↗

   Colab that introduces various aspects of the language and applies them to simple ML problems.

   
4. ▲
   Pax github.com

   tool category #jax link ↗

   A Jax-based machine learning framework for training large scale models.

5. ▲
   JAX: Accelerated Machine Learning Research | SciPy 2020 | VanderPlas youtube.com

   video category #jax link ↗

   JAX's core design, how it's powering new research, and how you can start using it.

   
6. ▲
   JAX GitHub Discussions github.com

   site category #jax link ↗

JAX › Tutorials and Blog Posts: “Describes the state of JAX and the JAX ecosystem at DeepMind.”