JAX
In plain wordsA free Google library that lets researchers write mathematics in familiar Python and automatically get derivatives and very fast code for graphics chips and Google's AI chips.
Google's open-source Python library for array computing on accelerators, built around a NumPyNumPy, prop. n.: The basic free Python library for working quickly with large grids of numbers; nearly every data-science and machine-learning tool in Python is built on it or copies its style.-like API and a small set of transformations that can be stacked on any numerical function: grad for automatic differentiationautomatic differentiation, n.: A technique that lets a computer work out exactly how a program's output would change if its inputs changed slightly, by tracking every arithmetic step. It lets AI software compute training updates automatically. (forward or reverse mode, nested to any depth), jit for just-in-time compilation with the XLAXLA, abbr.: A compiler from Google that turns AI programs into fast code for chips such as TPUs and GPUs, mainly by merging many small steps into fewer, bigger ones. compiler, vmap for automatic batching, and sharding primitives that spread arrays and computation over many GPUs or TPUsTPU, abbr.: Google's own family of chips built specifically for AI, used to train and run its Gemini models and rented to other companies through its cloud.. Transformations work by tracing pure functions into an intermediate representation, so JAX code is written in a functional style, with immutable arrays and explicitly passed random-number keys. Neural-network libraries such as Flax and Equinox, the optimizer library Optax and large-scale training codebases build on it. Google DeepMindGoogle DeepMind, prop. n.: Google's main artificial intelligence laboratory, formed in 2023 by joining DeepMind and Google Brain. It builds the Gemini models and scientific AI systems such as AlphaFold. uses JAX for much of its large-model training, and it is also popular in scientific computing.
Wrap the loss in jax.grad, vmap it over examples, and jit the whole update step.
Origin
Created at Google by Roy Frostig, Matthew James Johnson and Chris Leary, building on Autograd, a library developed at Harvard by Dougal Maclaurin, David Duvenaud, Matthew Johnson and Ryan Adams. It was described at the SysML conference in 2018 and open-sourced later that year.
Usage note
Not a full framework in the style of PyTorchPyTorch, prop. n.: The most widely used free software toolkit for building and training neural networks, especially in research; programmers write ordinary Python and it handles the calculus and the graphics chips. or KerasKeras, prop. n.: A free, beginner-friendly Python library for building neural networks from ready-made layers in a few lines of code, while a larger toolkit underneath does the heavy computing.: it supplies arrays and transformations, while higher-level libraries supply layers and training loops. Keras 3 can use JAX as a backend.
Related
- XLAXLA, abbr.: A compiler from Google that turns AI programs into fast code for chips such as TPUs and GPUs, mainly by merging many small steps into fewer, bigger ones.
- automatic differentiationautomatic differentiation, n.: A technique that lets a computer work out exactly how a program's output would change if its inputs changed slightly, by tracking every arithmetic step. It lets AI software compute training updates automatically.
- TPUTPU, abbr.: Google's own family of chips built specifically for AI, used to train and run its Gemini models and rented to other companies through its cloud.
- NumPyNumPy, prop. n.: The basic free Python library for working quickly with large grids of numbers; nearly every data-science and machine-learning tool in Python is built on it or copies its style.
- PyTorchPyTorch, prop. n.: The most widely used free software toolkit for building and training neural networks, especially in research; programmers write ordinary Python and it handles the calculus and the graphics chips.
- TensorFlowTensorFlow, prop. n.: Google's free software toolkit for building, training and running machine-learning models; it led the field after 2015 and still runs many models inside phones, browsers and company systems.
- KerasKeras, prop. n.: A free, beginner-friendly Python library for building neural networks from ready-made layers in a few lines of code, while a larger toolkit underneath does the heavy computing.
Sources
- Frostig, R., Johnson, M. J. & Leary, C. (2018). Compiling Machine Learning Programs via High-Level Tracing. SysML Conference 2018.
Mentioned in
AI stackAI stack, n.: The layers of technology an AI product depends on, from chips and data centers at the bottom, through the models, up to the apps people use, each built on the layer below., automatic differentiationautomatic differentiation, n.: A technique that lets a computer work out exactly how a program's output would change if its inputs changed slightly, by tracking every arithmetic step. It lets AI software compute training updates automatically., CUDACUDA, prop. n.: Nvidia's software system for programming its graphics chips to do general calculations. Most AI software is built on it, which is a big reason Nvidia chips dominate AI., KerasKeras, prop. n.: A free, beginner-friendly Python library for building neural networks from ready-made layers in a few lines of code, while a larger toolkit underneath does the heavy computing., NumPyNumPy, prop. n.: The basic free Python library for working quickly with large grids of numbers; nearly every data-science and machine-learning tool in Python is built on it or copies its style., PyTorchPyTorch, prop. n.: The most widely used free software toolkit for building and training neural networks, especially in research; programmers write ordinary Python and it handles the calculus and the graphics chips., tensortensor, n.: A block of numbers with any number of dimensions: a single number, a list, a grid, or a stack of grids. AI software stores and processes all its data this way., TensorFlowTensorFlow, prop. n.: Google's free software toolkit for building, training and running machine-learning models; it led the field after 2015 and still runs many models inside phones, browsers and company systems., TPUTPU, abbr.: Google's own family of chips built specifically for AI, used to train and run its Gemini models and rented to other companies through its cloud., LiteRTLiteRT, prop. n.: Google's free software for running AI models on phones, small gadgets and microcontrollers; it was called TensorFlow Lite until 2024 and is used inside many Android and iPhone apps., MLXMLX, prop. n.: A free Apple software toolkit for running and training AI models efficiently on Mac computers with Apple's own chips, popular for running language models locally on a laptop., MuJoCoMuJoCo, prop. n.: A free physics simulator, now maintained by Google DeepMind, that AI researchers use to train and test simulated robots and bodies, especially with trial-and-error learning., NCCLNCCL, prop. n.: NVIDIA's software library that lets its graphics chips exchange data with each other quickly — inside one computer and across networks — and that nearly all large-scale AI training on NVIDIA hardware relies on., ROCmROCm, prop. n.: AMD's free, open-source software for programming its graphics and AI chips — its answer to Nvidia's CUDA, letting AI software run on AMD hardware., XLAXLA, abbr.: A compiler from Google that turns AI programs into fast code for chips such as TPUs and GPUs, mainly by merging many small steps into fewer, bigger ones.