JAX: Accelerated Machine Learning Research via Composable Function Transformations in Python
summary
JAX gives NumPy-style code hardware acceleration by composing function transformations: grad, jit, vmap implement autodiff, end-to-end XLA compilation, and automatic batching. Demo shows a 5000×5000 float32 matmul running in 12 ms on a TPU, per-example gradients via grad+vmap+jit, and dense Hessians by composing forward/reverse modes. At scale JAX enters MLPerf with 1.5–2.5× faster submissions and reports extreme speedups (example: from days to seconds), including a 38.35 s WMT run on 1024 TPU cores (global batch 2048).
AI-written summary · about
chapters (7)
- 0:00Introduction to Jax
- 2:26Basic usage and performance
- 3:54Core function transformations
- 7:04Advanced composition
- 8:34Research applications
- 9:35Scaling and infrastructure
- 11:52Large-scale training demo
related talks
01h 10mJAX: 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