JAX: array computing#
JAX is the foundation of the JAX AI Stack! It provides high-performance array computing functionality over accelerators via a simple NumPy-like API and function transformations.
Main JAX site: jax.dev
If you’d like to learn more about JAX beyond what’s covered in the Getting started with JAX for ML guide, we recommend starting with the JAX tutorials.
The JAX docs cover many other useful topics including: