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

Since we are assuming you are fairly familiar with NumPy, this chapter is by no means a thorough coverage of the topic and will just serve as a quick refresher. We will practice using JAX while discovering the equivalence between normal NumPy and the JAX version of NumPy's syntax. 

>**Note:** As a reference, we'll use `np`  and `jnp` respectively for default and JAX NumPy versions in our codes. 

# Inner product

The inner product of two vectors can be calculated by any of the three syntaxes (**`dot()`**, **`inner()`** and **`@`**), as shown below:



a = jnp.arange(1,10)
b = jnp.arange(11,20)

print(jnp.dot(a,b))
print(jnp.inner(a,b))
print(a@b)
print(b@a)

# Linear functions

A linear function can be represented as:

$$ f(x) = \alpha x_1+\beta x_2+\gamma x_3+....
$$

We can compose it as a lambda function too:

a= 1.2
b = 2.4

FuncA = lambda x : a*x[0]+b*x[1]

x = jnp.arange(1,3)
print(FuncA(x))

# Taylor approximation

Taylor approximation lies at the core of major optimization algorithms in machine learning. We can first-order approximate the value of a function near a point _**z**_ as:

$$ \hat{f}(x) = f(z) +\nabla f(z)^T (x-z)$$




Having already seen the efficiency of `grad()`, we can easily implement the Taylor approximation. For example, if at $z = (2,3.2)$ we want to approximate the value of the following function: 

$$ f(x) = 2x_1+3x_2
$$


Learn how you can perform various vector operations using JAX.

Introduction

JAX Programming Model

Linear Algebra

Random Variables and Distributions

JAX Ecosystem

Project: GAN Using the JAX ecosystem

Appendix

Vectors

Inner product

Linear functions