AI Potluck
Infrastructure / Core ML frameworks & libraries

JAX

Google

JAX is a Python library for accelerator-oriented array computation and program transformation, combining NumPy-style arrays with composable function transforms (automatic differentiation, JIT compilation via XLA, and auto-vectorization). It is developed by the jax-ml organization (originating at Google) and is widely used for high-performance numerical computing and large-scale ML research.

Verified live 2026-06-22 via primary sources. Apache-2.0 licensed with full public source.

Openness

5 high confidence
5.0
license
Apache-2.0(OSI)
source
public
core-gated
ungated

Apache-2.0 licensed with full public source.

Adoption

5 high confidence
5.0

Re-scored 4->5. 18.92M PyPI downloads in the last 30 days. The map's prevailing usage_volume level-5 floor is >10M/mo (cf. pydantic-ai ~31M, llama-index ~10.18M, langgraph ~58M, all level 5); the prior level 4 applied a stricter 5M-50M band used only in one ml_frameworks batch, which sat below map peers of equal volume. JAX at ~19M/mo is the foundation of a major ML ecosystem (Flax, the Google research stack) and clears the prevailing floor.

Capability

5 high confidence
5.0

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

Unchanged since 2026-07-30 (last edited, not re-checked)