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


In the challenge, we'll continue using the notations from the previous lessons, such as $I$ for the input image and so on.

# Challenge

Let's try to master convolution arithmetic with some practice. In this task, we have $I$ of size $1024\times1024$ and will convert it into a $256\times256$ by a series of convolutions.

Since there are a number of ways to design the above convolution network, don't worry too much about following any specific pattern. 

Below is a calculator in case you need some help in designing the filters, padding, and so on.

>**Note: ** The calculator widget's cells work just like a typical spreadsheet software, such as Microsoft Excel. For example, to calculate the answer of 75+12, type **`=75+12`** and press "Enter".


Let's apply what we have learned about convolution using JAX in this coding exercise.

Introduction

JAX Programming Model

Linear Algebra

Random Variables and Distributions

JAX Ecosystem

Project: GAN Using the JAX ecosystem

Appendix

Challenge: Convolution