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

# Discrete distributions

There are some scenarios where we need to sample the distribution from discrete events like a coin toss result, the value of a dice roll, or the number of people.

# Probability Mass Function (PMF)

The concept of CDF is the same for both continuous and discrete distributions, though PDF is replaced by the **Probability Mass Function (PMF)**. Here we calculate the probability for a given sample directly instead of differentiation.

$$ p_X(x) = P(X = x) $$

# Poisson distribution

Poisson distribution is used to sample **unlikely events** (events with low probability).

$$ f(k; \lambda)= \Pr(X{=}k)= \frac{\lambda^k e^{-\lambda}}{k!} $$

Both expected value and mean will be the same:

$$ E[X] = Var[X] = \lambda$$

We can sample a Poisson distribution through ```poisson(key, <λ's value>, size) ```.

This lesson will provide an overview of discrete distributions in JAX.

Introduction

JAX Programming Model

Linear Algebra

Random Variables and Distributions

JAX Ecosystem

Project: GAN Using the JAX ecosystem

Appendix

Discrete Distributions

Discrete distributions

Probability Mass Function (PMF)