JAX is a numerical computing library for Python that looks like NumPy but can compile code and compute gradients automatically on GPUs and TPUs.
In plain words
If you know NumPy, JAX code looks familiar. Its power is in transformations: functions that take your Python function and return a new one. grad gives its gradient, which training needs. jit compiles it to run fast. vmap runs it over a whole batch at once. These can be combined, and the same code runs on CPU, GPU or TPU.
A quick example
Imagine you write a function that measures a model's error on one example. Wrapping it in grad gives a function for the gradient, vmap makes it work on a batch of examples, and jit compiles the result, all without rewriting the original function.
Where you'll meet it
- Large-scale machine learning research.
- Scientific computing that needs gradients.
Related
Sources
- Docsjax-ml/jax, JAX (GitHub) · read 29 Sept 2026
- DocsJAX documentation, JAX · read 29 Sept 2026
Checked by WhatIsAI against the sources above · 29 Sept 2026. Spotted an error? Report it.