JAX: accelerated machine learning research via composable function transformations in Python
summary
JAX claims Python+NumPy code runs on accelerators by tracing pure functions and applying composable transformations: jit via XLA compiles whole programs and fuses elementwise ops (demo 5k×5k matmul ≈12 ms vs CPU ~200 ms), vmap for per-example gradients vectorizes batch plumbing so one-line per-example grads and kernel gram matrices are fast, and pmap across 8 TPU cores shards arrays and runs eight 5k×5k multiplies in parallel. The system composes grad, jvp/vjp, vmap, jit, pmap (used in MLPerf scale runs) and is practical enough to build neural-ODE regularization and DP workflows.
AI-written summary · about
chapters (8)
- 0:00Start
- 2:06Motivating JAX
- 31:50Transforming and staging Python functions
- 33:57Step 1: Python function + JAX IR
- 39:13Step 2: transform jaxpr
- 41:38Why researchers like JAX
- 43:12Limitations
- 44:54MLPerf 2020 Results
related talks
14mJAX: Accelerated Machine Learning Research via Composable Function Transformations in Python
11mIntro to JAX: Accelerating Machine Learning research
31mHigh performance machine learning with JAX
02h 24mMagical NumPy with JAX
23mProbabilistic Programming with Vectorized Programmable Inference
26mTraining Large Deep Learning Models: PyTorch vs JAX
41mJAX for Bayes
08mDemo: JAX, Flax and Gemma