Learning JAX
2h 24mAdvanced2022-10-19
Authors
Janani Ravi
Certified Google Cloud Architect and Data Engineer
Course details
In this course, instructor Janani Ravi gives you an in-depth look at JAX, a new experimental Python library designed for high performance, scientific computing and machine learning. Janani takes you through all aspects of JAX and what it is capable of, including: just-in-time compilation; automatic vectorization and automatic parallelization; computing gradients; performing transformations on pytrees; training simple neural networks; and more.
Skills covered
Machine LearningAdvancedArtificial Intelligence (AI)Programming LanguagesSoftware Development
Concepts
0. Introduction
- 01 - Prerequisites
- 02 - What is JAX
- 03 - Why use JAX
- 04 - Choosing JAX
- 05 - JAX vs. TensorFlow vs. PyTorch
- 06 - Getting set up with Colab
1. Working with JAX Arrays
- 07 - JAX arrays
- 08 - JAX arrays and NumPy arrays - Similarities
- 09 - JAX arrays and NumPy arrays - Differences
- 10 - Asynchronous dispatch and JAX array speed up
2. Just-in-Time Compilation
- 11 - Composable function transformations
- 12 - JIT and pure functions
- 13 - Using JIT
- 14 - Tracer objects in JIT
- 15 - Impure functions and JIT - I O streams
- 16 - Impure functions and JIT - Global state
- 17 - Impure functions and JIT - Iterators
- 18 - Jaxprs
- 19 - Control flow statements and JIT
- 20 - Static arguments in jitted functions
- 21 - Lambdas and JIT
3. Automatic Vectorization and Automatic Parallelization
- 22 - Understanding vectorization and parallelization
- 23 - Automatic vectorization
- 24 - Comparing naive and manual batching with automatic vectorization
4. Computing Gradients
- 25 - Understanding gradient computation
- 26 - Gradient computation
- 27 - Higher order gradients
- 28 - Jacobians and Hessians
5. Performing Transformations on pytrees
- 29 - Understanding pytrees
- 30 - Simple pytrees
- 31 - Operations on pytrees
- 32 - Pytree containers
- 33 - Custom containers as pytrees
6. Training Simple Neural Networks
- 34 - Regression using a single neuron - Loading and preprocessing data
- 35 - Regression using a single neuron - Helper functions
- 36 - Regression using a single neuron - Training and evaluating a model
- 37 - Regression using a neural network - Helper functions
- 38 - Regression using a neural network - Training a model and visualizing results
- 39 - Classification using neural network - Loading and preprocessing data
- 40 - Classification using neural network - Training and evaluating model
Conclusion
- 41 - Summary and next steps