JAX Tutorial series

This is a work in progress lecture series that helps to get started with Google JAX.

• Lecture 1: Introduction to JAX.
• Lecture 2: JIT-compilation, Functional Programming and Gradients
• Lecture 3: Loops and Conditions in JAX
• Lecture 4: Vectorization and Parallelization in JAX
• Lecture 5: Linear and Logistic regression with JAX

References

This lecture series uses some of the examples and inspiration from the following sources

• QuantEcon's JAX series
• JAX Documentation