--- id: jax version: "0.11.0" license: Apache-2.0 license_treatment: permissive maintenance: active --- # jax — Differentiate, compile, and transform Numpy code. License: permissive · Maintenance: active · Popularity: top 1,000 on PyPI ## Install pip install jax uv add jax poetry add jax ## Description
logo
# Transformable numerical computing at scale [![Continuous integration](https://github.com/jax-ml/jax/actions/workflows/ci-build.yaml/badge.svg)](https://github.com/jax-ml/jax/actions/workflows/ci-build.yaml) [![PyPI version](https://img.shields.io/pypi/v/jax)](https://pypi.org/project/jax/) [**Transformations**](#transformations) | [**Scaling**](#scaling) | [**Install guide**](#installation) | [**Change logs**](https://docs.jax.dev/en/latest/changelog.html) | [**Reference docs**](https://docs.jax.dev/en/latest/) ## What is JAX? JAX is a Python library for accelerator-oriented array computation and program transformation, designed for high-performance numerical computing and large-scale machine learning. JAX can automatically differentiate native Python and NumPy functions. It can differentiate through loops, branches, recursion, and closures, and it can take derivatives of derivatives of derivatives. It supports reverse-mode differentiation (a.k.a. backpropagation) via [`jax.grad`](#automatic-differentiation-with-grad) as well as... ## AI interpretation — verify before relying JAX is a Python library for high-performance numerical computing that automatically differentiates NumPy code, compiles functions with XLA for GPUs/TPUs, and provides composable transformations like grad, jit, and vmap for machine learning at scale. Verdict: JAX is a mature, actively maintained library (Production/Stable) with low install friction and no known vulnerabilities. Its permissive Apache-2.0 license and broad hardware support make it suitable for research and production machine learning. The main consideration is the Python 3.12+ requirement and the learning curve around JAX's functional programming model and compilation constraints. [View on SkillFed](https://skillfed.io/packages/jax) · [View on PyPI](https://pypi.org/project/jax/)