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

Which options is **not** a parameter of **`vmap()`**?

Which options is not a parameter of <code>vmap()</code>?

Both are used for parallelization; there is no significant difference.

`vmap()` is used for device parallelization and `pmap()` is used for batch vectorization.

`pmap()` is used for device parallelization and `vmap()` is used for batch vectorization.

Both are used for vectorization; there is no significant difference.

What's the main difference between `vmap()` and `pmap()`?

What’s the main difference between <code>vmap()</code> and <code>pmap()</code>?

Both are used for second-order Taylor approximation.

Both are applied on functions $f: \mathcal R^n \to \mathcal R$.

What's the common feature between gradient and Hessian matrices?

What’s the common feature between gradient and Hessian matrices?

**`linalg.eigvals(Hessian_at_x) `** returns the following values:

>[0.+0.j 6.+0.j 0.+0.j 0.+0.j]

The matrix $x$ is a _______ matrix.

<code>linalg.eigvals(Hessian_at_x) </code> returns the following values:
<blockquote>
[0.+0.j 6.+0.j 0.+0.j 0.+0.j]
</blockquote>
The matrix <math><semantics><mrow><mi>x</mi></mrow><annotation encoding="application/x-tex">x</annotation></semantics></math>x is a _______ matrix.

Let's test our knowledge of multivariate calculus, auto vectorization, and autoparallelization.

Introduction

JAX Programming Model

Linear Algebra

Random Variables and Distributions

JAX Ecosystem

Project: GAN Using the JAX ecosystem

Appendix

Quiz: Multivariate Calculus, Auto Vectorization/Parallelization