Quickstart: How to think in JAX — JAX documentation

JAX is a library for array-oriented numerical computation (à la NumPy), with automatic differentiation and JIT compilation to enable high-performance machine learning research. This document provides a quick overview of essential JAX features, so you can get started with JAX: JAX provides a unified NumPy-like interface to computations that run on CPU, GPU, or TPU, in local or distributed settings. JAX features built-in Just-In-Time (JIT) compilation via Open XLA, an open-source machine learning compiler ecosystem. JAX functions support efficient evaluation of gradients via its automatic differentiation transformations. JAX functions can be automatically vectorized to efficiently map them over arrays representing batches of inputs. JAX can be installed for CPU on Linux, Windows, and macOS directly from the Python Package Index: or, for NVIDIA GPU: For more detailed platform-specific installation information, check out Installation. Key concepts: JAX provides a NumPy-inspired interface f

saved by

• Benedict Neo