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

While JAX has powerful features, coding our deep learning applications can still be tricky. This isn't surprising, since JAX is intended to be a generic numeric computational library.

JAX does offer some pretty useful libraries for designing neural networks, though. We'll review them in this chapter and consolidate our understanding to build the project at the conclusion.

# Flax

As a high-performance neural network library, **Flax** aims to provide flexible designs while coding in JAX.

The main packages in Flax are:

* Neural networks
* Utilities

# Neural networks

The package **`flax.linen`** is used for all the required neural network classes. Because of the wide range of functionalities for neural networks, we'll restrict ourselves to only the most relevant ones here.




The use of **`flax.nn`** [is obsolete](https://docs.google.com/document/d/1hYavTVPaKVVe9Be8pCB7yW7r6dDv3RALVNit8NZca4c/edit) and shouldn't be confused with **`flax.linen`**. 

# Linen philosophy

Linen designers acknowledge that PyTorch and Tensorflow are already established libraries in the deep learning community, so ___there is little need to build a clone of either___.

>**Note:** This section is more about Linen’s philosophy and can be skipped without any loss of continuity.

Flax Linen instead focuses on the strength of JAX $-$ beauty of working with the pure functions $-$ and allows us to:

* Use **`vmap()`** for auto-vectorization of batches.
* Handle ragged batches in NLP.
* Increase the efficiency of convolutional neural networks by utilizing rematerialized **` scan`**.
* Remove memory headaches.

## Interoperability with JAX transforms

JAX operates on pure functions, which means we cannot mutate the state of the input variable. This restriction can be troublesome while making a neural network since it can have tons of parameters changing in each iteration. In order to overcome this hindrance, Linen modules are designed to be **re-functionalized** $-$ they can be recast (_in-place_) as explicit functions:

$$f(v_{in},x)\to v_{out},y$$

"***Code shouldn't get in the way.***" One should remember that purely functional code is unhampered by the above feature. We can use Linen modules inside them without any modification.

## Parameters representation

Similar to other libraries, Flax uses the power of **Pytrees** for storing the (neural network) parameters.

This lesson will offer an overview of Flax, a dedicated neural network library of JAX.

Introduction

JAX Programming Model

Linear Algebra

Random Variables and Distributions

JAX Ecosystem

Project: GAN Using the JAX ecosystem

Appendix

Introduction to Flax and Linen

Flax

Neural networks

Linen philosophy