GO — Healthy across the board

• Last commit today
• 34+ active contributors
• Distributed ownership (top contributor 23% of recent commits)
• Apache-2.0 licensed
• CI configured
• Tests present

_{Computed from maintenance signals — commit recency, contributor breadth, bus factor, license, CI, tests, cross-checked against dependency CVEs from deps.dev and OpenSSF Scorecard}

JAX is a Python library for composable transformations of numerical computing programs, enabling automatic differentiation, JIT compilation via XLA, vectorization, and other functional transformations on NumPy-like code. It unifies AD (reverse/forward modes), compilation to GPU/TPU, and program transformation into a single, composable system that lets users write pure Python functions that run at scale on accelerators. Monorepo structure: jax/ contains the core Python API (transformations, dtypes, numpy compatibility), jax/_src/ holds implementation details (ad.py for autodiff, interpreters/ for lowering to XLA), jaxlib/ wraps C++ bindings and XLA compilation. Bazel (BUILD files throughout) manages dependencies and cross-platform builds. CI workflows in .github/workflows/ test on CPU, CUDA, ROCm, and TPU hardware independently.

Machine learning researchers, numerical computing engineers, and physics simulators who need to differentiate through complex Python code, compile numerical kernels to GPUs/TPUs, or apply multiple transformations (grad, vmap, jit) to the same functions without rewriting them.

Highly mature and actively developed. 21.4M lines of Python across a large codebase with extensive CI coverage (50+ workflows in .github/workflows/ for CPU/GPU/TPU/ROCm targets), robust Bazel build system, comprehensive testing via pytest/hypothesis, and documented sharp edges. This is production-grade but labeled a research project with intentional design boundaries.

Low risk for stability but moderate friction for contributors: large monolithic codebase requiring Bazel builds, heavy C++/XLA backend coupling (2.5M lines of C++), and strict transformation semantics that reject impure Python (side effects, mutable state). Dependency on XLA upstream and hardware-specific backends (CUDA, ROCm, TPU) adds deployment complexity. Single-library ownership by Google research team may slow external feature adoption.

Active development visible in extensive CI/CD setup with presubmit tests on multiple hardware platforms, nightly NumPy compatibility testing (numpy_nightly.yml, oldest_supported_numpy.yml), and specialized workflows for TPU (cloud-tpu-ci-*), CUDA (bazel_cuda_h100_b200.yml), and ROCm. Recent additions suggest focus on new hardware support (B200, Mosaic), array API standardization (jax-array-api.yml), and backward compatibility testing.

git clone https://github.com/jax-ml/jax.git
cd jax
pip install -e . --no-build-isolation
python -c "import jax; print(jax.numpy.array([1,2,3]))"

For development with Bazel (required for some components): bazel build //jax/... and bazel test //jax/...

Daily commands: JAX is a library, not a service. For examples: python -c "import jax; jax.jit(lambda x: x**2)(5)"; see jax/_src/nn/ for neural network utilities, or run examples from official documentation. No local server or daemon needed; distributed training requires explicit setup (pmap, xmap, or jax.experimental.mesh_pmap).

• jax/__init__.py — Main entry point exposing JAX's public API including transformations like grad, jit, vmap, and core array operations.
• jax/_src/interpreters/xla.py — Core XLA compilation and execution backend that translates JAX operations to hardware-accelerated code.
• jax/_src/api.py — Implements high-level transformation APIs (grad, jit, vmap, pmap) that are fundamental to JAX's value proposition.
• jax/_src/dispatch.py — Dispatch logic routing computations to appropriate backends (CPU, GPU, TPU) based on device placement.
• jax/_src/dtypes.py — Type system and dtype handling that ensures numerical precision across heterogeneous device architectures.
• setup.py — Build configuration defining JAX's compilation flags, platform detection, and jaxlib integration.

• Tracer & Core _(jax._src.core, jax.src.dtypes) — Records operations during function tracing and builds abstract computation graphs independent of concrete values.
  • Failure mode: Shape mismatch or unsupported operation during trace causes early failure before compilation.
• AD Interpreter _(jax._src.interpreters.ad, jax.src.api.grad) — Applies reverse-mode differentiation rules to traced operations to compute gradients.
  • Failure mode: Missing or incorrect vjp rule for a primitive results in undefined gradient or wrong result.
• XLA Lowering & Compiler _(jax._src.interpreters.xla, jax.src.interpreters.mlir) — Lowers JAX operations to MLIR/XLA, applies fusion and optimization, generates device-specific code.
  • Failure mode: Compilation timeout, out-of-memory, or XLA bug causes hang or crash on first jit call.
• Device Backend & Runtime _(jaxlib (C++), jax.src.dispatch, CUDA/HIP/TPU drivers) — Manages device buffers, executes compiled kernels, synchronizes results, and coordinates inter-device communication.
  • Failure mode: GPU out-of-memory, device disconnection, or kernel error causes runtime exception or silent wrong results.
• NumPy Compatibility Layer _(jax.src.numpy., jax._src.scipy.) — Implements NumPy-like API (linalg, random, fft, array ops) on JAX primitives with correct broadcasting and dtype handling.
  • Failure mode: Incorrect promotion rules or unsupported dtype combinations cause dtype errors or silent precision loss.
• Distributed/pmap Coordinator _(jax._src.pxla, jax.src.distributed) — Partitions functions across multiple devices, manages collective operations, and synchronizes execution.
  • Failure mode: Device deadlock, inconsistent sharding across processes, or collective op mismatch causes hang or deadlock.

• User Python code → JAX Tracer — User calls jax.grad(f)(x); JAX replaces x with a Tracer to record operations.
• Tracer operations → AD/XLA Interpreter — Recorded operations are passed to interpreters to build vjp computation graph or XLA lowering.
• Interpreter graph → XLA compiler — Computation graph is lowered to MLIR, optimized, and compiled to device-specific code.
• Compiled executable → Device backend — Binary is uploaded to GPU/TPU; input buffers are copied to device memory.
• Device execution result → JAX Array — Device buffer is wrapped in a JAX Array (lazy reference) and returned to user.

Add a new NumPy-compatible function

1. Define the primitive operation signature in jax/_src/interpreters/xla.py lower_builtin (e.g., register an XLA computation). (jax/_src/interpreters/xla.py)
2. Implement the function in the appropriate module (e.g., jax/_src/numpy/linalg.py or jax/_src/numpy/math.py). (jax/_src/numpy/linalg.py)
3. Register it in jax/_src/numpy/init.py to expose in the public jax.numpy namespace. (jax/_src/numpy/__init__.py)
4. Add gradient rules via custom_vjp in jax/_src/ad.py or jax/_src/custom_derivatives.py if differentiation is needed. (jax/_src/custom_derivatives.py)

Add a new transformation (like jit or grad)

1. Define the high-level API function in jax/_src/api.py (e.g., my_transform with wrapper and @curry decorator). (jax/_src/api.py)
2. Implement the interpreter logic in jax/_src/interpreters (e.g., a new ad.py-style module for your interpreter). (jax/_src/interpreters/ad.py)
3. Export it from jax/init.py to make it part of the public API. (jax/__init__.py)
4. Add tests in tests/ directory following pytest conventions and integration with hypothesis for property testing. (tests/)

Add support for a new primitive operation

1. Register the primitive in jax/_src/core.py using Primitive class and register_lowering for each backend. (jax/_src/core.py)
2. Implement XLA lowering in jax/_src/interpreters/xla.py lower_builtin function. (jax/_src/interpreters/xla.py)
3. Add AD rules in jax/_src/interpreters/ad.py using defvjp or defvjp_closure for backward differentiation. (jax/_src/interpreters/ad.py)
4. Optionally add forward-mode rules in jax/_src/api.py or interpreters/ad.py using defjvp. (jax/_src/api.py)

Add a new backend or device support

1. Add platform detection and initialization in jax/_src/lib/init.py or config.py. (jax/_src/config.py)
2. Implement backend-specific lowering rules in jax/_src/interpreters/xla.py for device placement and execution. (jax/_src/interpreters/xla.py)
3. Update setup.py to conditionally compile C++ bindings and link against platform-specific libraries (e.g., CUDA, ROCm). (setup.py)
4. Add CI workflow in .github/workflows/ (e.g., bazel_cuda.yml) to test compilation and correctness on the new platform. (.github/workflows/bazel_cuda.yml)

• XLA (Accelerated Linear Algebra) — Provides portable, optimized compilation to multiple hardware backends (CPU, GPU, TPU) while abstracting vendor-specific details.
• MLIR (Multi-Level Intermediate Representation) — Enables flexible, composable lowering from high-level JAX operations through multiple optimization levels to target-specific code.
• Python + NumPy API — Leverages existing NumPy ecosystem familiarity while enabling JIT compilation and automatic differentiation of pure Python functions.
• C++/pybind11 for jaxlib bindings — Achieves high-performance device communication and low-level runtime operations without Python GIL contention.
• PyTree abstraction — Handles arbitrary nested structures (dicts, lists, custom objects) uniformly for transformations, avoiding explicit flattening by users.

• Traced (functional) computation model over imperative state mutation

  • Why: Enables safe composition of transformations (grad∘jit∘vmap) and static analysis for optimization.
  • Consequence: Requires users to structure code as pure functions; control flow must use jnp.where or other functional primitives, not Python if/for.

• Delayed compilation with on-demand lowering to XLA

  • Why: Amortizes compilation cost across multiple calls and allows shape-polymorphic optimization.
  • Consequence: First call to a jitted function is slower; shape/dtype changes trigger recompilation, risking bloat in dynamic workloads.

• Stateless random number generation (via explicit PRNG keys)

  • Why: Ensures determinism, reproducibility, and safe parallelization across devices without hidden mutable state.
  • Consequence: API differs from NumPy; users must thread keys through code, adding verbosity.

• Reverse-mode AD (grad) as primary differentiation mechanism

  • Why: Efficient for scalar outputs and typical neural network backpropagation.
  • Consequence: Forward-mode (jvp) available but less optimized; inefficient for Jacobian-vector products with many inputs.

• JAX is not a machine learning framework itself (no built-in model zoo, optimizer, layers); frameworks like Flax/Equinox build on top.
• Does not provide automatic GPU memory management beyond XLA; users must understand device placement and buffer lifetimes.
• Not designed for imperative, stateful programming (mutable arrays, in-place operations); tracing model requires pure functions.
• Does not guarantee numerical stability (e.g., ill-conditioned matrices in grad); users responsible for numerical care.
• Not optimized for sparse tensors or irregular workloads; designed for dense, regular computations.

• In-place mutations in traced code — User code calling jitted functions; jax._src.core tracing: Modifying a list/dict

Immutability required: JAX rejects in-place mutations (e.g., x[i] = y); use functional updates instead. Tracing semantics: control flow (if/while) must be data-independent or use jax.lax.cond/scan; standard Python if-statements don't work under jit/vmap. XLA/Bazel required: C++ extensions need bazel build, not pure pip install for development. Hardware-specific tests: CUDA/ROCm tests need actual GPUs; CPU-only installs miss backend detection. Version pinning: jaxlib must match jax version closely; mismatches cause cryptic C++ linker errors.

• Automatic Differentiation (Reverse/Forward Mode) — Core JAX capability; understanding VJP (vector-Jacobian product) and JVP (Jacobian-vector product) is essential to use grad/jax.linearize correctly and diagnose AD errors
• Functional Transformations / Higher-Order Functions — JAX's design model — grad, jit, vmap, etc. are functions that return functions; requires thinking compositionally rather than imperative
• Program Tracing & Symbolic Execution — jit/vmap work by tracing your Python function on abstract values to build an XLA graph; causes counter-intuitive behavior (print statements disappear, control flow fails) if not understood
• Pytrees (Nested Data Structures) — JAX treats dicts, lists, tuples, and custom classes as tree-like data; pytree registration (tree_util.register_pytree_node) is how you make custom classes work with vmap/grad
• XLA Intermediate Representation (IR) — JAX lowering targets XLA IR (HLO); understanding that your Python code becomes a functional dataflow graph helps debug performance and understand compilation limits
• Pure Functions & Immutability — JAX strongly enforces functional purity (no side effects, no mutations); violating this causes silent wrong results or cryptic errors under transformations
• Collective Operations (all_reduce, pmap, xmap) — For distributed training and multi-device parallelism; pmap/xmap are higher-level abstractions over SPMD execution that require understanding data partitioning across devices

• google/jaxlib — Provides the compiled C++ backends and XLA bindings that JAX's Python API depends on; distributed separately but tightly coupled
• openxla/xla — The XLA compiler framework that JAX uses for lowering and optimization; JAX is the primary Python consumer of XLA
• pytorch/pytorch — Competing deep learning framework with similar auto-diff and JIT compilation; JAX and PyTorch users often compare or use both
• numpy/numpy — JAX's NumPy API is fully compatible by design; tracks NumPy's evolution and maintains alignment for user migration
• deepmind/dm-tree — Inspired JAX's pytree utilities; provides alternative tree traversal for users integrating both libraries

To work on one of these in Claude Code or Cursor, paste: Implement the "<title>" PR idea from CLAUDE.md, working through the checklist as the task list.

Add comprehensive CI workflow for JAX Array API standard compliance

JAX has a .github/workflows/jax-array-api.yml workflow, but there's no dedicated documentation or validation suite. The Array API standard is critical for JAX's interoperability. A new contributor could expand test coverage for Array API compliance by creating a structured test module that validates all Array API functions against the standard spec, with clear error messages for deviations.

• [ ] Review existing .github/workflows/jax-array-api.yml to understand current scope
• [ ] Create tests/array_api_compliance_test.py with systematic tests for each Array API function
• [ ] Add validation that JAX's Array API exports match the official standard namespace
• [ ] Document expected Array API behavior in a new section of CONTRIBUTING.md
• [ ] Run workflow locally with pytest tests/array_api_compliance_test.py to verify

Implement missing documentation for GPU/TPU-specific transformations and behaviors

The repo has extensive CI workflows for CUDA, ROCm, and TPU (.github/workflows/bazel_cuda.yml, bazel_rocm.yml, bazel_test_tpu.yml) but lacks a dedicated user guide explaining device-specific behavior of JAX transformations like jit, vmap, and pmap across different backends. This is a critical gap for users deploying to these devices.

• [ ] Create docs/device_specific_transformations.md documenting CUDA, ROCm, and TPU behavior
• [ ] Add examples showing how jit, pmap, and vmap differ across GPU/TPU backends
• [ ] Document memory layout implications and performance tips per device type
• [ ] Add cross-references in main documentation files pointing to this new guide
• [ ] Include troubleshooting section based on existing GitHub issues for device-specific problems

Add integration tests for Sphinx documentation build pipeline

The repo has complex Sphinx configuration (pydata-sphinx-theme==0.16.1, sphinx-book-theme==1.1.1, myst-nb>=1.0.0) with pinned versions to avoid breaking changes, but no automated validation that docs build successfully with all dependencies. A new contributor could create a lightweight documentation build test to catch regressions early.

• [ ] Create .github/workflows/docs-build-test.yml that runs sphinx-build on the docs directory
• [ ] Add pre-commit hook in .pre-commit-config.yaml to validate Sphinx syntax locally
• [ ] Create tests/docs_build_test.py that validates all RST/Markdown files parse correctly
• [ ] Document the build process in README.md under a 'Building Documentation' section
• [ ] Test against the pinned dependency versions from pyproject.toml to ensure compatibility

• Add comprehensive type stubs (.pyi files) for jax.numpy module — currently incomplete, blocking IDE autocomplete and type checkers used by downstream projects.
• Extend the pytree documentation and add more examples in docstrings of jax/_src/tree_util.py — pytree registration is powerful but underdocumented, causing confusion for new users handling custom classes.
• Create a 'transformation recipe' guide in docs/ showing common patterns (composing grad+jit, vmap over batch dimension, pmap for multi-host) with minimal examples — many users struggle with correct composition order.

The JAX repository demonstrates generally good security practices with CI/CD workflows and dependency management, but has several medium-risk

• Medium · Overly Permissive Sphinx Documentation Dependencies — build/nonfreethreading-requirements.txt and build/parallel_accelerator_execute.sh (documentation dependencies). The dependency file pins sphinx<8.0 and uses specific versions with known constraints. While not critical, sphinx-book-theme==1.1.1 depends on pydata-sphinx-theme v0.15 which may have unpatched vulnerabilities. Additionally, snowballstemmer<3.0.0 constraint suggests potential compatibility issues that could mask security updates. Fix: Regularly audit Sphinx ecosystem dependencies for CVEs. Consider using pip-audit or similar tools in CI/CD pipeline. Update to latest compatible versions of sphinx, sphinx-book-theme, and pydata-sphinx-theme as security patches become available.
• Medium · Third-party GitHub Action Dependencies — .github/workflows/ and .github/actions/. Multiple GitHub workflow files reference external actions (download-jax-cpu-wheels, download-jax-cuda-wheels, download-jax-rocm-wheels) without pinned versions visible in the file structure. This creates supply chain risk if action repositories are compromised or deleted. Fix: Pin all GitHub Actions to specific commit SHAs rather than branch names or version tags. Implement SLSA framework for artifact verification. Audit third-party actions regularly for vulnerabilities.
• Medium · Unencrypted Artifact Storage and Upload — .github/workflows/ and .github/actions/upload-test-artifacts/. The workflow uses artifact upload/download actions (.github/actions/upload-test-artifacts, .github/actions/download-jax-*-wheels) without explicit evidence of encryption in transit or at rest. Test artifacts and pre-built wheels may contain sensitive information. Fix: Ensure GitHub artifact storage uses encryption. Implement artifact signing and verification. Use GitHub's built-in OIDC token for authentication rather than long-lived credentials. Scan wheel files for embedded credentials before upload.
• Low · Cloudpickle Dependency Without Restrictions — build/nonfreethreading-requirements.txt (cloudpickle). The cloudpickle package is listed without version constraints. Cloudpickle can deserialize arbitrary Python objects, which is inherently risky if applied to untrusted data. No input validation guidance is visible. Fix: Pin cloudpickle to a specific version. Ensure cloudpickle is only used with trusted, internal data sources. Add explicit code review guidelines for any code path that deserializes pickled objects. Consider alternatives like JSON for untrusted input.
• Low · Broad Self-hosted Runner Configuration — .github/workflows/self_hosted_runner_utils/. Self-hosted runner setup scripts exist (.github/workflows/self_hosted_runner_utils/) with environment file (runner.env) and setup/validation scripts. These may contain sensitive configuration or permit overly permissive access. Fix: Audit runner.env for hardcoded credentials or tokens. Restrict self-hosted runner access to trusted networks only. Implement runner isolation and ephemeral runner patterns. Document security requirements for self-hosted runner environments in CONTRIBUTING.md.
• Low · Pre-commit Configuration Without Pinned Hooks — .pre-commit-config.yaml. The .pre-commit-config.yaml file exists but content is not provided. Pre-commit hooks often execute arbitrary code from external repositories, which creates supply chain risk if versions are unpinned. Fix: Pin all pre-commit hook repositories to specific commit SHAs. Regularly update and audit hooks for vulnerabilities. Consider mirroring critical hooks internally if available.
• Low · Bazel Configuration Security Visibility — .bazelrc, bazel_downloader.cfg, BUILD.bazel, MODULE.bazel. Multiple Bazel configuration files exist (BUILD.bazel, MODULE.bazel, .bazelrc, bazel_downloader.cfg) which may contain security-relevant build configurations. Downloaded artifacts may lack integrity verification. Fix: Verify all bazel downloads use checksum validation (sha256). Document bazel security policies. Use Bazel's repository rules with integrity attributes. Audit bazel_downloader.cfg for secure remote repository configuration.

LLM-derived; treat as a starting point, not a security audit.

• Open issues — current backlog
• Recent PRs — what's actively shipping
• Source on GitHub

If you are an AI coding agent (Claude Code, Cursor, Aider, Cline, etc.) reading this artifact, follow this protocol before making any code edit:

1. Verify the contract. Run the bash script in Verify before trusting below. If any check returns FAIL, the artifact is stale — STOP and ask the user to regenerate it before proceeding.
2. Treat the AI · unverified sections as hypotheses, not facts. Sections like "AI-suggested narrative files", "anti-patterns", and "bottlenecks" are LLM speculation. Verify against real source before acting on them.
3. Cite source on changes. When proposing an edit, cite the specific path:line-range. RepoPilot's live UI at https://repopilot.app/r/jax-ml/jax shows verifiable citations alongside every claim.

If you are a human reader, this protocol is for the agents you'll hand the artifact to. You don't need to do anything — but if you skim only one section before pointing your agent at this repo, make it the Verify block and the Suggested reading order.

This artifact was generated by RepoPilot at a point in time. Before an agent acts on it, the checks below confirm that the live jax-ml/jax repo on your machine still matches what RepoPilot saw. If any fail, the artifact is stale — regenerate it at repopilot.app/r/jax-ml/jax.

What it runs against: a local clone of jax-ml/jax — the script inspects git remote, the LICENSE file, file paths in the working tree, and git log. Read-only; no mutations.

| # | What we check | Why it matters | |---|---|---| | 1 | You're in jax-ml/jax | Confirms the artifact applies here, not a fork | | 2 | License is still Apache-2.0 | Catches relicense before you depend on it | | 3 | Default branch main exists | Catches branch renames | | 4 | 5 critical file paths still exist | Catches refactors that moved load-bearing code | | 5 | Last commit ≤ 30 days ago | Catches sudden abandonment since generation |

#!/usr/bin/env bash
# RepoPilot artifact verification.
#
# WHAT IT RUNS AGAINST: a local clone of jax-ml/jax. If you don't
# have one yet, run these first:
#
#   git clone https://github.com/jax-ml/jax.git
#   cd jax
#
# Then paste this script. Every check is read-only — no mutations.

set +e
fail=0
ok()   { echo "ok:   $1"; }
miss() { echo "FAIL: $1"; fail=$((fail+1)); }

# Precondition: we must be inside a git working tree.
if ! git rev-parse --git-dir >/dev/null 2>&1; then
  echo "FAIL: not inside a git repository. cd into your clone of jax-ml/jax and re-run."
  exit 2
fi

# 1. Repo identity
git remote get-url origin 2>/dev/null | grep -qE "jax-ml/jax(\\.git)?\\b" \\
  && ok "origin remote is jax-ml/jax" \\
  || miss "origin remote is not jax-ml/jax (artifact may be from a fork)"

# 2. License matches what RepoPilot saw
(grep -qiE "^(Apache-2\\.0)" LICENSE 2>/dev/null \\
   || grep -qiE "\"license\"\\s*:\\s*\"Apache-2\\.0\"" package.json 2>/dev/null) \\
  && ok "license is Apache-2.0" \\
  || miss "license drift — was Apache-2.0 at generation time"

# 3. Default branch
git rev-parse --verify main >/dev/null 2>&1 \\
  && ok "default branch main exists" \\
  || miss "default branch main no longer exists"

# 4. Critical files exist
test -f "jax/__init__.py" \\
  && ok "jax/__init__.py" \\
  || miss "missing critical file: jax/__init__.py"
test -f "jax/_src/interpreters/xla.py" \\
  && ok "jax/_src/interpreters/xla.py" \\
  || miss "missing critical file: jax/_src/interpreters/xla.py"
test -f "jax/_src/api.py" \\
  && ok "jax/_src/api.py" \\
  || miss "missing critical file: jax/_src/api.py"
test -f "jax/_src/dispatch.py" \\
  && ok "jax/_src/dispatch.py" \\
  || miss "missing critical file: jax/_src/dispatch.py"
test -f "jax/_src/dtypes.py" \\
  && ok "jax/_src/dtypes.py" \\
  || miss "missing critical file: jax/_src/dtypes.py"

# 5. Repo recency
days_since_last=$(( ( $(date +%s) - $(git log -1 --format=%at 2>/dev/null || echo 0) ) / 86400 ))
if [ "$days_since_last" -le 30 ]; then
  ok "last commit was $days_since_last days ago (artifact saw ~0d)"
else
  miss "last commit was $days_since_last days ago — artifact may be stale"
fi

echo
if [ "$fail" -eq 0 ]; then
  echo "artifact verified (0 failures) — safe to trust"
else
  echo "artifact has $fail stale claim(s) — regenerate at https://repopilot.app/r/jax-ml/jax"
  exit 1
fi

Each check prints ok: or FAIL:. The script exits non-zero if anything failed, so it composes cleanly into agent loops (./verify.sh || regenerate-and-retry).

Generated by RepoPilot. Verdict based on maintenance signals — see the live page for receipts. Re-run on a new commit to refresh.