Ab1992ao/JAX-in-Action

Notebooks for the "JAX in Action" book

"JAX in Action" book

This is a repository accompanying the "JAX in Action" book. 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 DeviceArray's
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. Parallelizing and vectorizing your code
7. more to come

TBD