AI Potluck
Contribute
Back to Gap Map Infrastructure / Core ML frameworks & libraries

JAX

Google
open source / Overall score: 4.4(strong)

JAX is a Python library for accelerator-oriented array computation and program transformation, combining NumPy-style arrays with composable function transforms: automatic differentiation, just-in-time compilation through XLA, and automatic vectorization. It is developed by the jax-ml organization, having originated at Google, and targets high-performance numerical computing and large-scale ML research.

Openness

5 high confidence
5.0
license
Apache-2.0(OSI)
source
public
core features withheld
no

Apache-2.0 licensed with full public source. The repository is public and unarchived, and the README describes the whole library with no paid, enterprise or hosted tier beside it, so source is public and the core ungated.

  • https://cdn.jsdelivr.net/gh/jax-ml/jax@main/LICENSE recorded 2026-09-29

    The LICENSE file in the jax-ml/jax repository is the full text of the "Apache License" "Version 2.0, January 2004", so the code is Apache-2.0 licensed.

  • https://cdn.jsdelivr.net/gh/jax-ml/jax@main/README.md recorded 2026-09-29

    The README describes JAX as "a Python library for accelerator-oriented array computation and program transformation" installed with pip. Searched the body for paid, enterprise, hosted, license key and self-host: none, so no gated tier sits beside the library.

  • https://docs.jax.dev/en/latest/developer.html recorded 2026-09-29

    The Building from source page tells readers to "First, obtain the JAX source code" with git clone from the jax-ml/jax repository, so the source is public.

  • https://github.com/jax-ml/jax recorded 2026-08-13

    Repo page still shows the Apache-2.0 license and the full public source tree

  • https://pypi.org/pypi/jax/json recorded 2026-09-29

    PyPI metadata for the jax package lists the license as "Apache-2.0" and the homepage as "https://github.com/jax-ml/jax".

Adoption

4 high confidence
4.0

Adoption is measured as PyPI downloads of the jax package in the trailing 30 days. JAX is also the foundation for other libraries such as Flax and the wider Google research stack, whose own installs are not counted here.

Capability

5 high confidence
5.0

Defines the functional/accelerated-compute layer for ML research.

  • https://cdn.jsdelivr.net/gh/jax-ml/jax@main/README.md recorded 2026-09-29

    The README lists automatic differentiation, compilation with XLA, and vectorization: "By composing `jax.vmap` with `jax.grad` and `jax.jit`, we can get efficient" per-example gradients, and scaling modes for "thousands of devices".

  • https://github.com/jax-ml/jax recorded 2026-08-13

    README still describes accelerator-oriented array computation and program transformation - autodiff, XLA JIT, vectorization

Verified 2026-09-29, except some scores, still awaiting re-confirmation