Magical NumPy with JAX
summary
Claim: write idiomatic NumPy and get autodiff, JIT and parallelism without learning a new array library. Demos use the NumPy‑compatible JAX API with vmap, lax.scan, prng keys, plus grad and jit to compose transforms; notebooks show loopless message‑passing, Gaussian random walks, deterministic RNG and exercises with solutions on Binder. One adoption note: a reported >100× speedup on a real model vs TensorFlow 1.x. Code runs and composes; exercises reveal the functional style and where state must be made explicit.
AI-written summary · about
chapters (21)
- 0:00Start
- 0:29The Format of the Tutorial
- 6:19What Exactly Is Jax
- 8:12Prerequisite Knowledge
- 9:30Jupiter Notebooks
- 11:19Exercises in Jupiter Notebooks
- 14:26Loopless Loops
- 14:48Semantics of Vmap
- 17:20Sum
- 18:12Change the in-Axis Argument of Vmap
- 26:33Automatic Differentiation
- 47:52Message Passing
- 52:42Loopless Loops Exercise 3
- 1:10:42Semantics of Lacks Scan
- 1:11:34Cumulative Sum
- 1:13:02Exercises
- 1:49:05Notebook on Deterministic Randomness
- 1:52:41Gaussian Random Block
- 1:56:42Deterministic Randomness
- 1:57:01Deterministic Randomness Notebook
- 2:17:01Compilation
related talks
01h 10mJAX: accelerated machine learning research via composable function transformations in Python
14mJAX: Accelerated Machine Learning Research via Composable Function Transformations in Python
11mIntro to JAX: Accelerating Machine Learning research
31mHigh performance machine learning with JAX
23mProbabilistic Programming with Vectorized Programmable Inference
41mJAX for Bayes
27mAPL-Inspired Techniques for Advanced NumPy
37mUsing Numba Effectively Today