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.
 LinuxLinks helps you discover the best free and open source software for Linux. Explore our extensive collection of recommended applications, including alternatives to popular proprietary software.We also publish software reviews, Linux tutorials, hardware reviews, free programming books and articles covering the wider open source community. Browse our Linux resources to discover more. Know a useful free and open source Linux application we haven't covered? Suggest a program for inclusion. Has your software been featured on LinuxLinks? Display one of our Featured on LinuxLinks badges on your website or GitHub README. |
Please read our Comment Policy before commenting.