jax

Differentiate, compile, and transform Numpy code.

jax Permissive license Apache-2.0 Active ★ 36,156 v0.11.0 released 2026-07-16

Install

jax on PyPI

pip

pip install jax

uv

uv add jax

poetry

poetry add jax

Package facts

License	Apache-2.0 (permissive)
Python support	supports the current Python release (>=3.12)
Install friction	low — pure-Python wheel
Runtime dependencies	5 — jaxlib, ml_dtypes, numpy, opt_einsum, scipy
Maintenance	actively maintained — 28 days since the last release
Last repo commit	2026-08-13
First released	2018-12-07
Popularity	one of the top 1,000 most-downloaded packages on PyPI (30-day window, as of 2026-08-13)
Known vulnerabilities	none known (OSV.dev, checked 2026-08-13)

Evidence: jax-0.11.0-py3-none-any.whl

Development Status :: 5 - Production/StableProgramming Language :: Python :: 3.12Programming Language :: Python :: 3.13Programming Language :: Python :: 3.14Programming Language :: Python :: Free Threading :: 3 - Stable

About jax

from the package's own PyPI description — quoted content, verbatim

Transformable numerical computing at scale

Continuous integration (image) PyPI version (image)

Transformations | Scaling | Install guide | Change logs | Reference docs

What is JAX?

JAX is a Python library for accelerator-oriented array computation and program transformation, designed for high-performance numerical computing and large-scale machine learning.

JAX can automatically differentiate native Python and NumPy functions. It can differentiate through loops, branches, recursion, and closures, and it can take derivatives of derivatives of derivatives. It supports reverse-mode differentiation (a.k.a. backpropagation) via jax.grad as well as...

Read as markdown · JSON record · Source repository · Homepage

AI interpretation — verify before relying

AI-generated interpretation of the package facts above; every digit, version, license, or vulnerability id it cites is grounded in the facts already shown on this page

JAX is a Python library for high-performance numerical computing that automatically differentiates NumPy code, compiles functions with XLA for GPUs/TPUs, and provides composable transformations like grad, jit, and vmap for machine learning at scale.

Low install friction with a pure-wheel distribution (jax-0.11.0-py3-none-any.whl). Actively maintained with a release 28 days old; requires Python 3.12+. Five runtime dependencies (jaxlib, ml_dtypes, numpy, opt_einsum, scipy) are standard numerical libraries.

Licensed under Apache-2.0 (permissive), allowing commercial use, modification, and distribution with minimal restrictions—suitable for most production and research contexts.

Usage

pip install jax
import jax
import jax.numpy as jnp

def loss(params, inputs, targets):
  preds = jnp.dot(inputs, params) + 0
  return jnp.sum((preds - targets)**2)

grad_loss = jax.jit(jax.grad(loss))

Requires Python 3.12 or later. GPU/TPU acceleration requires jaxlib compiled for your hardware; CPU-only installation is available but slower.

Verdict: JAX is a mature, actively maintained library (Production/Stable) with low install friction and no known vulnerabilities. Its permissive Apache-2.0 license and broad hardware support make it suitable for research and production machine learning. The main consideration is the Python 3.12+ requirement and the learning curve around JAX's functional programming model and compilation constraints.

Needs verification

Whether jaxlib installation complexity or platform-specific binary availability poses practical friction beyond the wheel distribution.
Performance characteristics and real-world compilation overhead for typical ML workloads.
Stability and maturity of AMD GPU and Windows experimental support mentioned in the description.

automatic differentiation numpygpu accelerated array computationjit compilation machine learningvectorization vmap batchingdifferentiable programming frameworkxla compiler numerical computingbackpropagation autodiff librarytpu gpu scaling distributed