High performance machine learning with JAX
summary
JAX is a Python front end to XLA that gives automatic differentiation and a drop‑in jax.numpy replacement so code often runs on GPUs or TPUs after swapping imports with minor caveats. It exposes functional transforms like grad, jit and vmap to trace, compile and auto‑vectorize functions; pmap shards work across devices and psum aggregates gradients across cores. TPUs rentable as v3‑8 present eight cores (pods scale to thousands), and user examples report these pipelines running blindingly fast.
AI-written summary · about
chapters (29)
- 0:00Introduction
- 1:20Welcome
- 1:40Overview
- 2:12What is JAX
- 3:18JAX
- 5:00Autograde
- 5:20Gradient
- 6:15Tracing
- 7:53Justintime compilation
- 8:06JAX jit
- 8:54JAX compilation
- 9:56Vectorization
- 11:36Vmap
- 12:59Augment example
- 14:17Batch example
- 14:59Batch dimensions
- 15:18TPU
- 16:03Xla
- 17:21Pmap example
- 19:06Collective operators
- 21:08Xmap
- 21:37JAX vs Keras
- 22:52JAX Libraries
- 25:29JAX Optics
- 26:02Update rule
- 28:11Recap
- 28:33Multihost
- 28:56Parallelization
- 29:24Summary
related talks
11mIntro to JAX: Accelerating Machine Learning research
14mJAX: Accelerated Machine Learning Research via Composable Function Transformations in Python
01h 10mJAX: accelerated machine learning research via composable function transformations in Python
02h 24mMagical NumPy with JAX
26mTraining Large Deep Learning Models: PyTorch vs JAX
41mJAX for Bayes
23mProbabilistic Programming with Vectorized Programmable Inference
14mAccelerating AI inference workloads