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

# Random variable

Consider a function:
$$f(x) = x^2 +5$$

No matter how many times we call this function, it will always give the same output for a given input.

On the other hand, consider some **random/stochastic function**; the seventh roll of the dice, the temperature at 3 PM, or the quantum state of a particle. All of them are never guaranteed to give the same output every time.

The outcome of this random function is known as a **random variable**. It's hard to think of a field that doesn't rely in some way on random variables. It's no wonder that they play a central role in modeling almost any real-world system.

# PRNG

One of the advantages computers have over manual work is their **deterministic** behavior: if we run an algorithm under the same conditions, it will always give the same result.

This advantage can sometimes turn out to be a disadvantage as well, though. Computer scientists have been aware of this trait for a long time and have thus devised algorithms known as **Pseudo-Random Number Generation (PRNG)**. 

John Von Neumann, who proposed one of the earliest PRNG algorithms, cautioned in his famous quote:
>"Anyone who considers arithmetical methods of producing random digits is, of course, in a state of sin." $-$ (1951)

## Mersenne Twister

One of the most commonly used PRNG, **Mersenne Twister** (named after the *Mersenne primes* it uses), is widely used in languages and libraries like C++, R, CUDA, NumPy, and more.

>Despite its broad implementation, Mersenne Twister has **[some issues](https://cs.stackexchange.com/a/53475)**.

## Threefry counter

Counter-based PRNGs, which we'll review shortly, address the shortcomings of Mersenne Twister and other conventional PRNGs. 
In 2011, D.E.Shaw Research introduced the **[Threefry counter](http://www.thesalmons.org/john/random123/papers/random123sc11.pdf)**, which is the baseline behind the JAX implementation.

We'll review the JAX implementation (referred to going forward as **JAX PRNG**) features.

---





















This lesson will introduce pseudo-random number generation in JAX and discuss how it offers an edge over other PRNGs.


Introduction

JAX Programming Model

Linear Algebra

Random Variables and Distributions

JAX Ecosystem

Project: GAN Using the JAX ecosystem

Appendix

Pseudo-Random Number Generation

Random variable