13 Resources

Turbocharge ML with JAX and TPUs

Welcome
1 TPUs
2 Introduction to JAX
3 Linear Regression in JAX
4 Exercise 1: MNIST in JAX
5 Flax Foundations
6 Exercise 2: Linear Regression in Flax
7 Let’s build a ResNet!
8 Diffusion for the curious
9 Stable Diffusion in JAX / Flax !
10 Vision Transformer
11 Exercise 1 Solution
12 Exercise 2 Solution
13 Resources
References

13 Resources

HuggingFace diffusers library

Google Research VIT

JAX on GKE Example (GPU)

12 Exercise 2 Solution