# jax-ml/jax

Composable transformations of Python+NumPy programs: differentiate, vectorize, JIT to GPU/TPU, and more

Repository: https://github.com/jax-ml/jax
Canonical: https://ross.abutalabs.com/products/jax
Homepage: https://docs.jax.dev
Language: Python
License: Apache-2.0
License Family: permissive
Topics: jax
Last push: 2026-08-27T00:18:16+00:00

## Health v2 (maintenance only)
Score: 99/100 (v2, computed 2026-09-03T02:20:16.233290+00:00)
- activity 99, release rhythm 98, longevity 100
- inputs: {"age_days": 2869, "days_push": 7, "days_rel": 16, "gap_med": 28.0, "n_releases_24m": 31}
- flags: none
- formula: round(0.45*activity + 0.35*rhythm + 0.20*longevity); archived -> min(score, 10)

## Adoption (not part of the score)
Stars 36215, forks 3744 (observed 2026-08-28T04:12:00.999641+00:00)

## What it is
JAX is a Python library for accelerator-oriented array computation and composable program transformations, offering NumPy-compatible APIs with automatic differentiation, vectorization, and JIT compilation via XLA. It scales the same code across CPU, GPU, and TPU backends and includes Pallas for writing custom hardware kernels.

## Use cases
- train neural networks with automatic differentiation in Python
- JIT-compile NumPy-style code to run on GPUs and TPUs
- vectorize per-example gradient computations with vmap
- write custom GPU/TPU kernels with Pallas
- run large-scale machine learning research across accelerators
- compute higher-order derivatives like Jacobians and Hessians
- parallelize and shard computations across multiple devices

## When to choose
- you need high-performance numerical computing or ML research on GPU/TPU
- you want composable autodiff, batching, and compilation in pure Python
- you prefer a NumPy-like API with accelerator scaling
- you need custom low-level kernels for TPU or Hopper+ GPUs

## When to avoid
- you need a batteries-included deep learning framework with prebuilt layers and training loops out of the box
- your project depends on impure Python code with side effects inside compiled functions
- you need stable experimental APIs like Pallas Triton or unsupported hardware backends
- you only need simple CPU array math where plain NumPy suffices

## Facets
- artifact type: library
- maturity: stable
- function: machine-learning, deep-learning, math, compiler, gpu-computing, llm-training
- domain: machine-learning, deep-learning, gpu-computing, data-science
- platform: python, windows, cloud
- tags: autodiff, jit-compilation, xla, numpy, tpu, vectorization, pallas, numerical-computing, algorithms, linux, macos, gpu

## Member repositories
- jax-ml/jax (main) score 99

## Provenance
- Observed fields: from GitHub, fetched 2026-08-28T04:12:00.999641+00:00.
- Health v2: computed from the inputs above; adoption is never an input.
- Inferred fields (summary, facets, guidance): AI-extracted, prompt v1, taxonomy v1, on 2026-08-29T16:26:43.094812+00:00, confidence not recorded.
  - readme: https://github.com/jax-ml/jax (fetched 2026-08-28T04:12:00.999641+00:00, sha 7c101dab4b0f)
  - homepage: https://docs.jax.dev (fetched 2026-08-29T07:47:46.993762+00:00, sha 7a112cd10fa7)
  - site_page: https://docs.jax.dev/en/latest/pallas/quickstart.html (fetched 2026-08-29T07:47:47.018800+00:00, sha 392db64cd25d)
  - site_page: https://docs.jax.dev/en/latest/pallas/CHANGELOG.html (fetched 2026-08-29T07:47:47.021460+00:00, sha 1ac6612ed34f)
  - site_page: https://docs.jax.dev/en/latest/about.html (fetched 2026-08-29T07:47:47.024717+00:00, sha dce08330eb76)
  - site_page: https://docs.jax.dev/en/latest/faq.html (fetched 2026-08-29T07:47:47.026494+00:00, sha c4cb7ecaba18)
  - site_page: https://docs.jax.dev/en/latest/installation.html (fetched 2026-08-29T07:47:47.012178+00:00, sha 271ba08ebebd)
  - site_page: https://docs.jax.dev/en/latest/pallas/tpu/quickstart.html (fetched 2026-08-29T07:47:47.014732+00:00, sha 76c783c3a037)
  - site_page: https://docs.jax.dev/en/latest/pallas/gpu/quickstart.html (fetched 2026-08-29T07:47:47.016791+00:00, sha bdc9e6b9b200)
  - site_page: https://docs.jax.dev/en/latest/changelog.html (fetched 2026-08-29T07:47:47.028550+00:00, sha b7dcbe7aa32c)
- Data as of 2026-08-30T08:39:29.467469+00:00.