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



Installing **JAX** is ridiculously easy, similar to **NumPy** or any other Python library.

# CPU-only installation

If you don't have GPU support, you can still code JAX on your CPU.

1. Upgrade the **`pip`**: **`pip install --upgrade pip`**
2. Install JAX: **`pip install --upgrade "jax[cpu]"`**





pip install --upgrade pip
pip install --upgrade "jax[cpu]"

---

# GPU installation

Like other deep learning systems, JAX includes support for CUDA, though we have to install the **CUDA** and **CuDNN** ourselves.

## Installing **CUDA**

Since CUDA depends on a number of parameters $-$ like OS family, distribution, version, and so on $-$ we are using the most common settings of **Ubuntu 20.04** on **`x86_64`** here:

Let's install JAX on the local machine and get working!

Introduction

JAX Programming Model

Linear Algebra

Random Variables and Distributions

JAX Ecosystem

Project: GAN Using the JAX ecosystem

Appendix

Installation

CPU-only installation

GPU installation

Installing CUDA