jax has 21.6M monthly downloads and 36K GitHub stars. Version 0.10.2 supports Python 3.11+, first published in 2018.
It is actively maintained under a permissive license, last released in Jun 2026.
It is built on top of jaxlib, ml-dtypes and numpy. It is used by 1,897 tracked packages, including orbax-checkpoint, flax and tensorflowjs.
GitHub topics include jax.
Package Insights
((week_daily_avg - month_daily_avg) / month_daily_avg) * 100Weekly Downloads
GitHub Stars
Downloads by OS
Python Versions
Top Countries
Dependencies
- jaxlib <=0.10.2,>=0.10.1
- ml-dtypes >=0.5.0
- numpy >=2.0
- opt-einsum
- scipy >=1.14
19 optional dependencies
- jax-cuda12-plugin[cuda12]
- jax-cuda12-plugin[cuda]
- jax-cuda12-plugin[cuda12-local]
- jax-cuda13-plugin[cuda13]
- jax-cuda13-plugin[cuda13-local]
- jax-rocm7-plugin[rocm7-local]
- jaxlib[minimum-jaxlib]
- jaxlib[ci]
- jaxlib[cuda]
- jaxlib[cuda12]
- jaxlib[cuda12-local]
- jaxlib[cuda13]
- jaxlib[cuda13-local]
- jaxlib[rocm7-local]
- jaxlib[tpu]
- kubernetes[k8s]
- libtpu[tpu]
- requests[tpu]
- xprof[xprof]
Used By
- orbax-checkpoint
- flax
- tensorflowjs
- optax
- chex
- equinox
- numpyro
- clu
- orbax-export
- e3nn-jax
- jax-md
- nequix
- jaxopt
- qwix
- aqtp
- pathwaysutils
- lineax
- tokamax
- seqio
- optimistix
- drjax
- seqio-nightly
- diffrax
- distrax
- blackjax
- rax
- klujax
- kaggle-environments
- gdsfactoryplus
- tpu-inference
- sax
- augmax
- mujoco-mjx
- jax-dataclasses
- jaxlie
- jaxellip
- dopamine-rl
- litert-torch
- gymnax
- ott-jax
- alphafold-colabfold
- tree-math
- brax
- interpax
- jax-solar
- rlax
- xdsl-jax
- google-metrax
- kauldron
- jax-triton