Why JAX?

JAX is a Python library designed for high-performance numerical computing, especially machine learning research. Its API for numerical functions is based on NumPy, a collection of functions used in scientific computing. Both Python and NumPy are widely used and familiar, making JAX simple, flexible, and easy to adopt.

In addition to its NumPy API, JAX includes an extensible system of composable function transformations that help support machine learning research, including:

• Differentiation: Gradient-based optimisation is fundamental to ML. JAX natively supports both forward and reverse mode automatic differentiation of arbitrary numerical functions, via function transformations such as grad, hessian, jacfwd and jacrev.
• Vectorisation: In ML research we often apply a single function to lots of data, e.g. calculating the loss across a batch or evaluating per-example gradients for differentially private learning. JAX provides automatic vectorisation via the vmap transformation that simplifies this form of programming. For example, researchers don’t need to reason about batching when implementing new algorithms. JAX also supports large scale data parallelism via the related pmap transformation, elegantly distributing data that is too large for the memory of a single accelerator.
• JIT-compilation: XLA is used to just-in-time (JIT)-compile and execute JAX programs on GPU and Cloud TPU accelerators. JIT-compilation, together with JAX’s NumPy-consistent API, allows researchers with no previous experience in high-performance computing to easily scale to one or many accelerators.

We have found that JAX has enabled rapid experimentation with novel algorithms and architectures and it now underpins many of our recent publications. To learn more please consider joining our JAX Roundtable, Wednesday December 9th 7:00pm GMT, at the NeurIPS virtual conference.

JAX at DeepMind

Supporting state-of-the-art AI research means balancing rapid prototyping and quick iteration with the ability to deploy experiments at a scale traditionally associated with production systems. What makes these kinds of projects particularly challenging is that the research landscape evolves rapidly and is difficult to forecast. At any point, a new research breakthrough may, and regularly does, change the trajectory and requirements of entire teams. Within this ever-changing landscape, a core responsibility of our engineering team is to make sure that the lessons learned and the code written for one research project is reused effectively in the next.

One approach that has proven successful is modularisation: we extract the most important and critical building blocks developed in each research project into well tested and efficient components. This empowers researchers to focus on their research while also benefiting from code reuse, bug fixes and performance improvements in the algorithmic ingredients implemented by our core libraries. We’ve also found that it’s important to make sure that each library has a clearly defined scope and to ensure that they’re interoperable but independent. Incremental buy-in, the ability to pick and choose features without being locked into others, is critical to providing maximum flexibility for researchers and always supporting them in choosing the right tool for the job.

Other considerations that have gone into the development of our JAX Ecosystem include making sure that it remains consistent (where possible) with the design of our existing TensorFlow libraries (e.g. Sonnet and TRFL). We’ve also aimed to build components that (where relevant) match their underlying mathematics as closely as possible, to be self-descriptive and minimise mental hops “from paper to code”. Finally, we’ve chosen to open source our libraries to facilitate sharing of research outputs and to encourage the broader community to explore the JAX Ecosystem.

Our Ecosystem today

Source: https://deepmind.com/blog/article/using-jax-to-accelerate-our-research