jax-ml/jax

JAX is a Python library from Google for high-performance numerical computation, especially designed for machine learning research. The core problem it solves is that writing fast, hardware-accelerated numerical code in Python, the kind that runs on GPUs and TPUs rather than just a regular CPU, normally requires significant knowledge of low-level tools. JAX lets you write familiar NumPy-style array code (NumPy is the standard Python library for mathematics on arrays) and then apply transformations that automatically make it faster, differentiable, or parallelizable. The four key transformations are: jax.grad, which automatically computes derivatives (gradients) of any function, essential for training machine learning models; jax.jit, which compiles a function using XLA (a hardware compiler from Google) so it runs much faster on CPUs, GPUs, or TPUs; jax.vmap, which vectorizes a function so it operates efficiently on batches of data at once; and sharding tools for distributing computation across hundreds or thousands of devices for large-scale training. Gradients can be taken through loops, conditions, and recursion, and you can take the gradient of a gradient of a gradient to any depth. These transformations compose freely, you can compile a vectorized gradient function with a single line of code. JAX is a research project rather than a user-facing product, and the README explicitly warns about "sharp edges", surprising behaviors in certain cases. It is used by researchers building custom machine learning models, scientific simulations, and large-scale training systems. The tech stack is Python, with installation via pip, supporting Linux, macOS, and Windows with CPU, GPU, and TPU backends.