JAX

JAX is an open-source Python library for high-performance numerical computing and machine learning, developed by Google Research. It provides a NumPy-like API (via jax.numpy) with automatic differentiation and relies on the XLA compiler to generate optimized code for CPU, GPU, and TPU backends. JAX emphasizes composable function transformations that enable gradients, vectorization, and parallelism across devices.

Its core features include automatic differentiation (reverse-mode for gradients, with facilities for higher-order derivatives), and specialized Jacobian and Hessian tools (grad, jacfwd, jacrev, hessian). It also offers vectorization through vmap, which maps a function over batch dimensions without explicit Python loops, and parallelization through pmap for data-parallel execution on multiple devices. The jit transform compiles Python functions with XLA, producing fast kernels and reducing Python overhead. JAX also provides a stateless random-number API via jax.random, using explicit PRNG keys.

JAX exposes a NumPy-like API through jax.numpy, while the lax module provides lower-level primitives for more control. The library encourages a functional programming style with pure functions and explicit state passing, which makes transformations reliable and composable. It is commonly used as a backend for neural network libraries such as Flax, Haiku, and Equinox, and serves as a platform for experimentation in research and production.

JAX runs on CPU, GPUs, and TPUs and supports multi-device computation. While powerful, it has a learning curve, and certain Python constructs and in-place mutations are not compatible with its transformations. Code often requires structuring to enable JIT, grad, vmap, and pmap without runtime errors.