jax

JAX implements a complete NumPy-compatible API (jax.numpy) that wraps lower-level LAX primitives, enabling users to write familiar NumPy code while maintaining full traceability for automatic differentiation. The implementation maps NumPy operations to JAX's intermediate representation (Jaxpr) through a tracer system that intercepts Python operations, building a computational graph without requiring explicit graph construction syntax. This allows seamless gradient computation and other transformations on NumPy-style code.

JAX implements automatic differentiation through a tracer-based interpreter system (jax.interpreters.ad) that builds a Jaxpr representation of a function, then applies reverse-mode (backpropagation) or forward-mode differentiation rules to compute gradients. The system supports higher-order derivatives (grad of grad), arbitrary nesting of AD with other transformations, and custom VJP/JVP rules for user-defined operations. Gradients are computed by tracing through the function once to build the computational graph, then applying chain rule transformations.

JAX implements a comprehensive type system (jax.dtypes) that handles numeric types (int32, float32, complex64, etc.) with automatic promotion rules. The system supports weak type promotion (e.g., Python int to int32) and strong type promotion (e.g., int32 to float32 in mixed operations). Type information is preserved through transformations and used by the compiler for optimization. Users can control promotion behavior via jax.numpy.promote_types and explicit casting.

JAX integrates with Google's XLA compiler by lowering Jaxpr intermediate representations to MLIR (Multi-Level Intermediate Representation) and StableHLO (Stable High-Level Operations). The lowering process converts high-level JAX operations to hardware-independent HLO, which XLA then optimizes and compiles to target-specific code (LLVM for CPU, NVPTX for GPU, HLO for TPU). This architecture enables single-source deployment across heterogeneous hardware without code changes.

JAX provides jax2tf and tf2jax bridges enabling seamless interoperability with TensorFlow. jax2tf converts JAX functions to TensorFlow SavedModel format, enabling deployment in TensorFlow ecosystems. tf2jax wraps TensorFlow operations as JAX functions, allowing mixed JAX/TensorFlow code. The bridges handle dtype conversion, device placement, and gradient flow, enabling gradual migration between frameworks or hybrid workflows.

JAX provides a configuration system (jax.config) enabling runtime control of behavior without code changes. Users can configure JIT defaults, device placement, dtype promotion, debugging flags, and experimental features. Configuration can be set via environment variables, Python API, or context managers, enabling flexible control of JAX behavior for different use cases (development, testing, production).

JAX's jit decorator traces a Python function to produce a Jaxpr intermediate representation, lowers it to MLIR/StableHLO, and compiles via XLA to hardware-specific executables (LLVM for CPU, NVPTX for GPU, HLO for TPU). The compilation pipeline exposes three stages (Traced, Lowered, Compiled) via jax.stages, allowing inspection and debugging of the compilation process. JIT compilation caches compiled functions by input shape and dtype, enabling fast re-execution of the same computation with different data.

JAX's vmap (vectorized map) transformation automatically vectorizes functions across a batch dimension by tracing the function once and generating SIMD/batched operations. Instead of writing explicit loops over batch dimensions, users annotate which axis to vectorize, and vmap generates efficient batched code that runs on vector units or tensor cores. The implementation uses a batching interpreter that transforms scalar operations into batched equivalents, composing with JIT for compiled vectorized kernels.

JAX's pmap (parallel map) distributes a function across multiple devices (GPUs, TPUs) by tracing the function and automatically generating sharded computation graphs. Each device receives a slice of the input along a specified axis, executes the function independently, and results are gathered. The system handles device placement, communication (all-reduce, all-gather), and synchronization transparently. pmap composes with JIT and grad, enabling distributed training and inference without explicit communication code.

JAX's Pallas subsystem enables writing custom kernels in a Python-like DSL that compiles to TPU (Mosaic) or GPU (Triton/Mosaic GPU) code. Pallas provides a lower-level abstraction than JAX's high-level operations, allowing fine-grained control over memory layout, communication patterns, and hardware-specific optimizations. Kernels written in Pallas integrate seamlessly with JAX's transformation system (grad, vmap, jit), enabling custom operations with automatic differentiation support.

JAX provides lax.cond, lax.while_loop, lax.for_loop, and lax.scan primitives that enable control flow within jitted and differentiated code. These primitives are implemented as special traced operations that build a Jaxpr representation of the control flow, enabling automatic differentiation and JIT compilation. Unlike Python's if/while statements, which are not traceable, lax primitives produce functional control flow that can be optimized and compiled.

JAX implements a stateless random number generation system (jax.random) that uses explicit seed/key management instead of global state. The system provides a threefry/philox counter-based PRNG that is deterministic, reproducible, and composable with transformations (vmap, pmap, jit). Keys are split and threaded through code, enabling parallel RNG streams without synchronization. The design avoids global state, making random operations safe in jitted and distributed code.

JAX's distributed computing system (jax.experimental.pjit and sharding annotations) enables automatic data and model parallelism across multiple devices. Users annotate arrays with sharding specifications (e.g., PartitionSpec), and the compiler automatically generates communication code (all-reduce, all-gather, reduce-scatter) to synchronize computations. The system integrates with XLA's collective operations and handles device placement transparently, enabling distributed training and inference without explicit communication code.

JAX's core architecture is built on composable function transformations (grad, jit, vmap, pmap) that operate on a pure functional intermediate representation called Jaxpr. Each transformation traces a function to produce a Jaxpr, applies interpretation rules (AD rules, batching rules, etc.), and produces a new function. Transformations can be arbitrarily nested and composed without special-casing, enabling powerful abstractions like grad(jit(vmap(f))) or vmap(grad(f)). The Jaxpr representation is a directed acyclic graph of primitive operations with explicit data flow.