jax-ml/jax

Composable transformations of Python+NumPy programs: differentiate, vectorize, JIT to GPU/TPU, and more

Structural Verdict

How Data Flows Through the System

Python functions with NumPy arrays flow through tracers that record operations into Jaxpr intermediate representation, then get lowered to MLIR and compiled to XLA for execution on accelerators

1. Function Tracing — Python function called with Tracer inputs to record operation graph
2. Jaxpr Construction — Traced operations converted to Jaxpr intermediate representation with typed variables
3. MLIR Lowering — Jaxpr lowered to MLIR HLO dialect with shape/type information
4. XLA Compilation — MLIR compiled to optimized XLA computation for target hardware
5. Execution — Compiled function executed on CPU/GPU/TPU with actual array data

System Behavior

How the system actually operates at runtime — where data accumulates, what loops, what waits, and what controls what.

Data Pools

Feedback Loops

• Recompilation Trigger (cache-invalidation, balancing) — Trigger: Function called with new shapes/dtypes. Action: Retrace and recompile function. Exit: Matching signature found in cache.
• Gradient Computation (recursive, balancing) — Trigger: grad() transformation applied. Action: Trace forward pass then construct reverse computation. Exit: All derivatives computed.

Delays & Async Processing

• JIT Compilation (async-processing, ~seconds) — First call to jitted function has compilation overhead
• Device Transfer (async-processing, ~milliseconds) — Data transfer between host and accelerator memory

Control Points

• jax_enable_x64 (env-var) — Controls: Whether to use 64-bit precision by default. Default: False
• jax_platform_name (env-var) — Controls: Which backend to use (cpu/gpu/tpu)
• jax_jit_pjit_api_merge (feature-flag) — Controls: Whether to merge jit and pjit APIs

Technology Stack

Key Components

• core (module) — Central tracing system with Tracer, Jaxpr, and abstract values for program transformation jax/_src/core.py
• jit (function) — Just-in-time compilation transformation that compiles functions to XLA jax/_src/api.py
• grad (function) — Automatic differentiation transformation for computing gradients jax/_src/api.py
• vmap (function) — Vectorization transformation that batches operations over array axes jax/_src/api.py
• MLIRLoweringContext (class) — Manages lowering of JAX operations to MLIR for XLA compilation jax/_src/interpreters/mlir.py
• PallasCall (class) — Custom kernel system for GPU programming with block-level operations jax/_src/pallas/core.py
• ffi_call (function) — Foreign function interface for calling external C/C++/CUDA code jax/_src/ffi.py
• ShapedArray (class) — Abstract value representing arrays with known shape and dtype jax/_src/abstract_arrays.py

Sub-Modules

Configuration

docs/autodidax2_part1.py (python-dataclass)

• primal (float, unknown)
• tangent (float, unknown)

docs/autodidax2_part1.py (python-dataclass)

• interpreter (Interpreter, unknown)
• primal (float, unknown)
• tangent (float, unknown)

docs/the-training-cookbook.py (python-dataclass)

• seq_length (int, unknown) — default: 128

jax/_src/api.py (python-dataclass)

• idx (int, unknown)
• primal (bool, unknown)

Science Pipeline

1. Array Creation — jnp.array() or device_put() converts Python data to JAX arrays [arbitrary Python nested structure → JAX ShapedArray with inferred shape/dtype] jax/_src/array.py
2. Function Tracing — Replace arrays with Tracers to record computation graph [concrete arrays → abstract ShapedArray values] jax/_src/core.py
3. Shape Inference — Propagate abstract shapes through operations without computing values [abstract shapes → inferred output shapes] jax/_src/abstract_arrays.py
4. Backend Lowering — Convert JAX operations to MLIR HLO with explicit tensor shapes [JAX abstract arrays → MLIR tensor types] jax/_src/interpreters/mlir.py

Assumptions & Constraints

• [warning] Assumes contracting dimensions have matching sizes but only checks at runtime during execution (shape)
• [info] Assumes total size is preserved during reshape but allows -1 inference with potential size mismatches (shape)
• [warning] Broadcasting assumes compatible dimensions but shape errors only surface during tracing (shape)

Related Ml Training Repositories

Frequently Asked Questions

Analyzed on March 31, 2026 by CodeSea. Written by Karolina Sarna.