Intro to JAX: Accelerating Machine Learning research
summary
JAX traces Python numerical functions into a jaxpr intermediate representation, applies composable transforms like grad, jit and vmap to automatically differentiate, compile to XLA HLO, vectorize and parallelize, then run the result on CPU, GPU or TPU with no code changes. It reports MLPerf wins on some benchmarks, powers large projects (AlphaFold runs on JAX), and supports pmap for multi-core/TPU parallelism while keeping a NumPy-like API for rapid prototyping.
AI-written summary · about
chapters (8)
- 0:00Intro
- 0:56JAX is Fast: MLPerf vo.7 Results SSD
- 1:39Deep learning in Numpy
- 2:59Motivating JAX
- 5:35JAX traces Python functions. What does this function do?
- 6:53Python function - JAX Intermediate Representation
- 8:21JAX is designed from ground-up around XLA
- 9:51JAX ecosystem JAX provides a foundation for a growing ecosystem of domain-specific tools: High-level Deep Learning Libraries Probabilistic Programming
related talks
01h 10mJAX: accelerated machine learning research via composable function transformations in Python
31mHigh performance machine learning with JAX
14mJAX: Accelerated Machine Learning Research via Composable Function Transformations in Python
02h 24mMagical NumPy with JAX
23mProbabilistic Programming with Vectorized Programmable Inference
41mJAX for Bayes
26mTraining Large Deep Learning Models: PyTorch vs JAX
08mDemo: JAX, Flax and Gemma