Machine Learning with Flax - From Zero to Hero

Intro - Flax is performant and reproducible
Deepnote walk-through sponsored
Flax basics
Flax vs Haiku
Benchmarking Flax
Linear regression toy example
Introducing Optax Adam state example
Creating custom models
self.param example
self.variable example
Handling dropout, BatchNorm, etc.
CNN on MNIST example
TrainState source code
CNN dropout modification
Outro and summary