"Deep Learning with JAX" book

This is a repository accompanying the "Deep Learning with JAX" book (originally called "JAX in Action"). Feel free to use the promotional code au35sap for 35% discount for this book and any other Manning's product.

Notebooks for the chapters:

1. Intro to JAX
   • JAX Speedup
2. Your first program in JAX
   • MNIST image classification with MLP in pure JAX
3. Working with tensors
   • Image Processing with Tensors
   • Working with jax.Array
   • Working with jax.Array (TPU version)
4. Autodiff
   • Different ways of getting derivatives
   • Working with gradients in TensorFlow, PyTorch, and JAX
   • Differentiating in JAX
5. Compiling your code
   • JIT compilation and more: JIT, Jaxpr, XLA, AOT
6. Vectorizing your code
   • Different ways to vectorize a function, Controlling vmap() behavior, More real-life cases
7. Parallelizing your computations
   • Using pmap()
   • Multi-host example
8. Parallelization using tensor sharding
   • Tensor sharding
9. Random numbers in JAX
   • Random augmentations, NumPy and JAX PRNGs
10. Complex structures in JAX/Pytrees
    • Pytrees, jax.tree_util functions, custom nodes
11. Higher-level neural network libraries
    • Using Flax: Simple MLP
    • Using Flax: Simple MLP with TrainState
    • Using Flax: Adding metrics
    • Flax: ResNet example with model state
    • HuggingFace Transformers: Using GPT-like models
    • HuggingFace Diffusers: Generating images with Stable Diffusion
12. Appendix: Experimental parallelization
    • Using xmap()
    • Using pjit()

Google AI/ML Developer Programs team supported this work by providing Google Cloud Credit.