jax-ml /

JAX

Live map
View repository
Snapshotmain · 2014e4c

How NumPy-like Python is traced, transformed, lowered, compiled, sharded, and dispatched to accelerators.

Click a building or route
Choose a path

Trace Python, apply layouts, lower the program, compile it, and launch on devices.

System map7 components · 8 connections
130%
JAX architecturePython APIs and primitives begin at left. Jaxpr transformations occupy the center, then sharding, lowering, compilation, and device execution proceed to the right.Python callable and ab…mesh and partition spe…partitioned StableHLO …compiled executable an…AutodiffinterpreterTransformationAPIPrimitivesand JaxprLoweringand compilerDevicedispatchBatchinginterpreterShardingplanner
Connections
BuildingsEntryServiceComputeStorageRuntimeToolingExternal

Terms, in plain English

TracerA stand-in value that records primitive operations during a transformation.

How to study this repo

  1. 1
    Print one Jaxpr

    Use make_jaxpr on a small function and match each equation to its Primitive definition.

  2. 2
    Compose transformations

    Compare grad(vmap(f)) with vmap(grad(f)), then follow both through staging and compilation.