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

We can consolidate our learning so far in a working example. 

Since this course slightly favors machine or deep learning, we'll practice with the ultimate 101 example: linear regression. In this challenge, you'll build linear regression step-by-step.

# Step 1: Regressor

Remember, linear regression fits the data point $x$ to a straight line:

$$ \hat{y}  = wx+b$$

As a starter, this part simply requires us to implement the regressor, $\hat{y}$




def Regressor(W,b):
  def y_(x):
    #Please update the line below with correct equation for y_
    return 1
  return y_

import jax
import jax.numpy as jnp
jax.config.update('jax_platform_name', 'cpu')

def Regressor(W,b):
  def y_(x):
    return W@x+b
  return y_

# Step 2: Loss function

The loss function for a single sample is:

$$\mathcal l(y_{'},y) = \frac{1}{2}(y^{'}-y)^2$$













We'll test the concepts learned so far by applying them in a linear regression example.

Introduction

JAX Programming Model

Linear Algebra

Random Variables and Distributions

JAX Ecosystem

Project: GAN Using the JAX ecosystem

Appendix

Challenge II: Linear Regression

Step 1: Regressor

Step 2: Loss function