JAX is a numerical computing and machine learning library that supports high-performance array operations, automatic differentiation, compilation, and accelerator execution. It is commonly used for research and model development requiring flexible, composable transformations of numerical programs.
An array computing library, bringing together XLA (Accelerated Linear Algebra) and automatic differentiation for high-performance numerical computing. JAX provides a simple and powerful API for writing accelerated numerical code with composable transformations. JAX provides features such as: - (automatic differentiation) - (just-in-time compilation) - (automatic vectorization or batching) - (parallelization) JAX is a language for expressing and composing transformations of numerical code, analogous---but much larger in scope---to Python's NumPy library. (In fact, the .numpy library under JAX is a functionally equivalent, but entirely rewritten version of the Python NumPy library.) JAX is particularly well-suited for speeding up many machine learning tasks by transforming the models and data into a form suitable for parallelism across GPU and TPU accelerator chips. Flax, Optax, Pax, and many other libraries are built on the JAX infrastructure.