Discover the power of JAX in deep learning. Gain insights into its ecosystem and learn about linear algebra, pseudo-random number generation, and optimization algorithms for cleaner, structured coding.

Doxx.tar.gz

NormalCode

JuPyter

JAX is a Python library designed for high-performance ML research. It is a powerful numerical computing library, just like Numpy, but with some key improvements.

In this course, you will learn all about JAX and its ecosystem of libraries (Haiku, Jraph, Chex, Flax, Optax). Addressing a wide range of audiences, you will cover several topics including linear algebra, random variables theory, pseudo-random number generation, and optimization algorithms.

By the end of this course, you will have a new set of skills that will make deep learning programming more intuitive, structured, and clean.

Introduction to JAX and Deep Learning

# Jacobian

The gradient is used for real-valued functions,  $f:R^n \to R$.

The concept of gradients can be extended to vector-valued functions, $f: R^n \to R^m$ by **Jacobian** matrices.

A Jacobian matrix is defined as:

$$ J = \begin{bmatrix}
    \dfrac{\partial \mathbf{f}}{\partial x_1} & \cdots & \dfrac{\partial \mathbf{f}}{\partial x_n} \end{bmatrix}
= \begin{bmatrix}
    \nabla^\mathsf{T} f_1 \\  
    \vdots \\
    \nabla^\mathsf{T} f_m   
    \end{bmatrix}
= \begin{bmatrix}
    \dfrac{\partial f_1}{\partial x_1} & \cdots & \dfrac{\partial f_1}{\partial x_n}\\
    \vdots & \ddots & \vdots\\
    \dfrac{\partial f_m}{\partial x_1} & \cdots & \dfrac{\partial f_m}{\partial x_n} \end{bmatrix} $$

JAX provides both **forward** and **reverse** mode auto-differentiation functions to calculate the Jacobian.

1. **`jacfwd()`** - calculates the Jacobian column-by-column
2. **`jacrev()`** - calculates the Jacobian row-by-row



The appendix complements the vector calculus lesson from the linear algebra chapter with a bit more details.

Introduction

JAX Programming Model

Linear Algebra

Random Variables and Distributions

JAX Ecosystem

Project: GAN Using the JAX ecosystem

Appendix

Vector Calculus - II

Jacobian