jax — JAX is a Python library for accelerator-oriented array computation and program transformation, designed for high-performance numerical computing and large-scale machine learning.

JAX is a Python library for accelerator-oriented array computation and program transformation, designed for high-performance numerical computing and large-scale machine learning.

About this project

JAX provides a system for composable function transformations of Python and NumPy programs. It is designed to scale numerical computing across hardware accelerators including GPUs and TPUs using XLA (Accelerated Linear Algebra) for compilation. Key capabilities include: - Automatic Differentiation: Using `jax.grad`, it supports reverse-mode and forward-mode differentiation through Python loops, branches, recursion, and closures, allowing for higher-order derivatives. - Just-In-Time (JIT) Compilation: Using `jax.jit`, it compiles pure functions via XLA to improve execution performance on accelerators. - Auto-vectorization: Using `jax.vmap`, it maps functions along array axes by pushing loops down to primitive operations, eliminating the need for manual batch dimension management. - Scaling and Parallelization: Supports various scaling modes including compiler-based automatic parallelization, explicit sharding, and manual per-device programming with explicit collectives. JAX is compatible with multiple platforms including Linux, macOS, and Windows, with support for CPU, NVIDIA GPU, Google TPU, AMD GPU, and experimental support for Apple and Intel GPUs.