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 |
| JAX | Python library for high-performance numerical computing |
| SageMath | Computer algebra system |
| SymPy | Library for symbolic mathematics |
| PyMC | Bayesian statistical modeling and probabilistic programming |
| Pyomo | Object-oriented algebraic modeling language |
| patsy | Package for describing statistical models and to build design matrices |
| mpmath | Library for arbitrary-precision floating-point arithmetic |
| SfePy | Finite element software package |
Read our verdict in the software roundup.
| Deep Learning with Python | |
|---|---|
| TensorFlow | A very popular Deep Learning framework |
| PyTorch | Tensors and Dynamic neural networks in Python |
| Keras | High-level neural networks API |
| fastai | Simplifies training fast and accurate neural networks |
| tinygrad | Compact, lightweight deep learning framework |
| PyTorch Lightning | Framework for scaling PyTorch training workflows |
| JAX | High-performance numerical computing and ML library |
| PyTensor | Library for fast numerical computation |
| MindSpore | Framework for training and deploying neural networks |
| Elephas | Distributed deep learning with Keras and Spark |
| Chainer | Powerful, flexible, and intuitive framework for neural networks |
| Equinox | Neural network library for building flexible models with JAX |
| Caffe | Convolutional Architecture for Fast Feature Embedding |
| TFlearn | Deep learning library featuring a higher-level API for TensorFlow |
| MXNet | Flexible and efficient library |
| CNTK | Distributed deep learning |
| Neupy | Python library for neural networks and deep learning |
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. |
Please read our Comment Policy before commenting.