jax: Transform and Accelerate Python Numerical Programs

Overview

JAX provides tools for writing numerical computations in Python and transforming them for differentiation, compilation, and vectorization. It addresses the challenge of taking familiar array-based code and making it suitable for high-performance execution on CPUs and accelerators, including GPUs and TPUs.

It is aimed at researchers and developers building scientific computing or machine learning workloads who need fine-grained control over how functions are transformed and scaled. JAX is a research project, and its documentation notes that it has sharp edges.

Key Features

• Compute gradients with jax.grad, including higher-order derivatives.
• Compile functions with jax.jit using XLA.
• Vectorize functions over array dimensions with jax.vmap.
• Compose differentiation, compilation, and vectorization.
• Scale computations using automatic, explicit, or manual sharding approaches.
• Run numerical workloads on supported CPU and accelerator platforms.

Use Cases

• Machine learning researchers can calculate model gradients and per-example gradients for training or analysis.
• Scientific computing developers can compile numerical functions to improve execution on supported hardware.
• Teams working with large accelerator workloads can shard computation across devices using JAX's parallel programming modes.
• Python developers with array-based programs can vectorize computations without manually adding batch dimensions throughout their code.

Project Facts

• Language: Python
• License: Apache-2.0
• Stars: 36.4k
• Forks: 3.8k
• Topics: jax
• Archived: no

Getting Started

Install the CPU package with:

pip install -U jax

For accelerator installation options and platform requirements, see the installation guide. For API and usage details, see the README and reference documentation.

Alternatives

• cupy: CuPy offers NumPy- and SciPy-compatible GPU arrays, while JAX adds program transformations such as automatic differentiation, compilation, and vectorization.

Considerations

• JAX transformations work best with functions that fit its programming model. In particular, jax.jit constrains the Python control flow a function can use.
• Accelerator support depends on the platform and installation path. The README lists some platform combinations as experimental or unsupported.
• The project describes itself as a research project and cautions users to expect sharp edges. Consult the gotchas guide before relying on assumptions from ordinary NumPy or Python behavior.