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

# Introductory example

We have already talked about the JIT's usefulness. Now let's take a look at the difference in speed between a JIT and a normal function:

#Normal Function
def FuncA(x,y):
  return np.asarray(np.dot(x,y))

#And JIT version
@jit
def FuncB(x,y):
  return jnp.dot(x,y)

x = jnp.ones((1000,1000))

start_time = timeit.default_timer()
FuncA(x,x)

print("Total time taken for normal function:", timeit.default_timer() - start_time)

start_time = timeit.default_timer()
FuncB(x,x)

print("Total time taken for JIT function:", timeit.default_timer() - start_time)

If we run the above code on the CPU, the output should be fairly straightforward. 

If we run it on GPU/TPU, on the other hand, the output is quite astonishing. The **JAX JIT** function is almost three times faster than the "non-JIT" function and completes within 0.02sec for a $1000\times1000$ array.

This result can be misleading, though. We need to examine some of the background mechanisms to fully understand what is happening here.

# Actual mechanism
When an operation (like `jnp.dot()` above) is executed, JAX doesn't wait for the whole result and instead returns a **`DeviceArray`** value. This may result in an erroneous result (well, more of a misleading result), but it can still be helpful. For example:

* We can pass a computation to another computation without waiting for the final result.
* We can infer the type or size of a **`DeviceArray`** in advance. 

It is also helpful that we can pipeline our operations without waiting for the outputs and thus can avoid being blocked.

However, if we want to determine the actual execution time, we should block the function call to wait until the computation is completed. The `block_until_ready()` operation does just this:

This lesson will introduce asynchronous dispatch.

Introduction

JAX Programming Model

Linear Algebra

Random Variables and Distributions

JAX Ecosystem

Project: GAN Using the JAX ecosystem

Appendix

Asynchronous Dispatch

Introductory example