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

# Probability distribution

A probability distribution is a function that gives the **probability of occurrence of a random event**. It can be either **discrete** or **continuous**, depending on the domain. We can formally define it as:

$$\mathcal P: \mathcal A\to \mathcal R$$

Here $\mathcal A$ can determine the distribution as either discrete or continuous as being a discrete or continuous set, respectively. 

There are a couple of useful ways to describe a probability distribution:

* Cumulative Distribution Function (CDF)
* Probability Density Function (PDF)

## Cumulative Distribution Function (CDF)

The CDF of a distribution is the value it will take for $x$ less than (or equal to) the given value (generalized as $X$). Formally,

$$F_X(x) = {P}(X\leq x) = \int_{-\infty}^x f_X(u) \, du$$

>**Note:** Since probabilities are always non-negative, CDF is a non-decreasing function.

## Probability Density Function (PDF)

More often, we are interested in calculating the value for a particular $x$ rather than the whole limit. Here we can use the Probability Density Function (PDF). This can be calculated by taking the derivative of the CDF:

$$f_X(x) = \frac{d}{dx} F_X(x)$$

# Moments

The expected value of a random variable is an intuitive measure of its **mean**, which is formally defined as:

$$\mathrm{E}[X] = \sum_{i=1}^\infty x_i\, p_i $$

As a simple example, suppose the probability of a dice roll is as follows:

$$P(1) = P(6) = 0.1$$
$$P(2) = P(3)=P(4) =P(5) = 0.2$$

The expected value will be:

$$\mathrm{E}[X] = \sum_{i=1}^6 x_i\, p_i =0.1(1)+0.2(2)+0.2(3)+0.2(4)+0.2(5)+0.1(6) =3.5$$

This notion of expected value can be extended to variance, etc., which are collectively known as **moments**.

The $\mathrm{n}^{th}$ moment can be defined as:


$$ \mu_n = \mathrm{E} \left[ ( X - \mathrm{E}[X] )^n \right]  = \int_{-\infty}^{+\infty} (x - \mu)^n f(x)\,\mathrm{d} x $$

The first four moments are:

1. **Mean**
1. **Variance**
1. Skewness
1. Kurtosis

Usually, we consider only mean and variance.

Now we'll revise some commonly used probability distributions along with their respective JAX functions.

---


This lesson will introduce the common probability distributions implementations in JAX.

Introduction

JAX Programming Model

Linear Algebra

Random Variables and Distributions

JAX Ecosystem

Project: GAN Using the JAX ecosystem

Appendix

Continuous Probability Distributions

Probability distribution

Cumulative Distribution Function (CDF)

Probability Density Function (PDF)

Moments