JAX is a Python library for high-performance numerical computing and large-scale machine learning.
It combines a NumPy-like API with composable program transformations, letting developers differentiate, compile, vectorize, and scale numerical code across CPUs, GPUs, TPUs, and other accelerators. The project is designed for accelerator-oriented array computation, research workflows, and machine learning systems that need automatic differentiation and efficient execution.
This is free and open source software.
Key Features
- Provides automatic differentiation for native Python and NumPy-style functions.
- Supports reverse-mode and forward-mode differentiation, including higher-order derivatives.
- Uses XLA to compile numerical programs for high-performance execution on supported hardware accelerators.
- Offers just-in-time compilation with jax.jit for pure Python functions.
- Includes auto-vectorization with jax.vmap to map functions over array axes efficiently.
- Supports scaling computations across multiple devices with automatic, explicit, and manual sharding approaches.
- Includes a NumPy-compatible array API through jax.numpy.
Website: github.com/jax-ml/jax
Support:
Developer: jax-ml
License: Apache License 2.0
JAX is written in Python. Learn Python with our recommended free books and free tutorials.
Related Software
| Python Mathematics Tools | |
|---|---|
| scikit-learn | Machine learning library for Python |
| NumPy | Core package for scientific computing with Python |
| SciPy | Ecosystem for mathematics, science, and engineering. |
| statsmodels | Statistical modeling and econometrics |
| SymPy | Library for symbolic mathematics |
| SageMath | Computer algebra system |
| patsy | Package for describing statistical models and to build design matrices |
| mpmath | Library for arbitrary-precision floating-point arithmetic |
Read our verdict in the software roundup.
 Explore our comprehensive directory of recommended free and open source software. Our carefully curated collection spans every major software category.This directory is part of our ongoing series of informative articles for Linux enthusiasts. It features hundreds of detailed reviews, along with open source alternatives to proprietary solutions from major corporations such as Google, Microsoft, Apple, Adobe, IBM, Cisco, Oracle, and Autodesk. You’ll also find interesting projects to try, hardware coverage, free programming books and tutorials, and much more. Discovered a useful open source Linux program that we haven’t covered yet? Let us know by completing this form. |