Need NumPy-style numerical code that can differentiate, compile, batch, and run across accelerators? JAX provides a functional-first Python library for high-performance array computing in machine learning, scientific computing, optimization, and simulation. section 1.0 section 1.2
JAX is useful when you need to:
- Compute with a NumPy-compatible multidimensional Array API while targeting CPUs, GPUs, or TPUs. section 1.1.1 section 1.2.1
- Compose automatic differentiation, just-in-time compilation, auto-vectorization, and parallel execution as function transformations. section 1.2.2
- Keep model parameters, random-number state, gradients, and optimizer state explicit rather than hidden inside objects. section 1.1.2
- Assemble a custom stack from separate neural-network, optimizer, data-loading, reinforcement-learning, or meta-learning modules. section 1.1.3
A functional array update
NumPy-style in-place updates do not fit JAX's immutable arrays. Instead of arr[i] += 10, write arr = arr.at[i].add(10). The same functional style applies to larger computations: define a pure function, then transform it with grad(), jit(), or vmap(). section 1.2.1 section 1.2.2
Figure 1.1 A simple (and pretty useless) function with some amount of computations that we will evaluate with pure NumPy, with JAX using a CPU and GPU, and with JAX-compiled versions of the same function for CPU and GPU backends. We will evaluate the time to calculate f(x).
Compare pure NumPy, JAX on CPU or GPU, and JIT-compiled JAX versions of the same computation to see where hardware and XLA compilation affect execution time. section 1.1.1
Figure 1.3 JAX (right) is much more than NumPy (left). It provides a set of composable function transformations for Python + NumPy code for compilation, vectorization, parallelization, and automatic differentiation.
Use the transformation overview to see how JAX extends NumPy with compilation, vectorization, parallelization, and automatic differentiation. section 1.2.2
The tradeoff is a change in habits: side effects, hidden state, mutable arrays, and some ordinary if statements or loops can conflict with JIT-oriented code. JAX is deliberately minimal compared with PyTorch and TensorFlow, so its ecosystem supplies many separate components. Existing embedded deployments or large established codebases may be better left in their current frameworks. section 1.1.2 section 1.3 section 1.1
