JAX (software)

JAX is a Pythonlibrary for accelerator-oriented array computation and program transformation, designed for high-performance numerical computing and large-scale machine learning. It is developed by Google with contributions from Nvidia and other community contributors.^[1][2][3]

It is described as bringing together a modified version of the automatic differentiation system autograd^[4] and OpenXLA's XLA (Accelerated Linear Algebra). It is designed to follow the structure and workflow of NumPy as closely as possible and works with various existing frameworks such as TensorFlow and PyTorch.^[5][6] The primary features of JAX are:^[7]

1. Providing a unified NumPy-like interface to computations that run on CPU, GPU, or TPU, in local or distributed settings.
2. Built-in Just-In-Time (JIT) compilation via OpenXLA, an open-source machine learning compiler ecosystem.
3. Efficient evaluation of gradients via its automatic differentiation transformations.
4. Automatic vectorization to efficiently map functions over arrays representing batches of inputs.

• Flax ^[8]
• Equinox ^[9]
• Optax ^[10]
• Diffrax

• NumPy
• TensorFlow
• PyTorch
• CUDA
• Accelerated Linear Algebra
• Comparison of machine learning software
• List of numerical libraries

• Documentationː docs.jax.dev
• Colab (Jupyter/iPython) Quickstart Guideː colab.research.google.com/github/google/jax/blob/main/docs/notebooks/thinking_in_jax.ipynb
• TensorFlow's XLAː www.tensorflow.org/xla (Accelerated Linear Algebra)
• YouTube TensorFlow Channel "Intro to JAX: Accelerating Machine Learning research": www.youtube.com/watch?v=WdTeDXsOSj4
• Original paperː mlsys.org/Conferences/doc/2018/146.pdf