0% completed

All LessonsFree Lessons (6)

Introduction

Course Introduction JAX Overview

JAX Programming Model

Pure Functions JAX and NumPy Just-in-Time (JIT) Compilation JAX Expressions (Jaxpr)Asynchronous Dispatch Auto-differentiation Pytrees Quiz: Basics of JAX Challenge: Basics of JAX

Linear Algebra

Vectors Matrices More on Matrices Quiz - Vectors and Matrices Challenge: Vectors and Matrices Auto Vectorization Auto Parallelization Multivariate Calculus Quiz: Multivariate Calculus, Auto Vectorization/Parallelization Challenge II: Linear Regression Convolutions Two-dimensional Convolutions Transposed Convolutions Quiz: Convolution Challenge: Convolution

Random Variables and Distributions

Pseudo-Random Number Generation Continuous Probability Distributions Discrete Distributions Quiz: PRNG and Probability Distributions Challenge: PRNG and Probability Distribution Divergence Measures Introduction to Optimal Transport and OTT

JAX Ecosystem

Introduction to JAX Ecosystem

Introduction to Flax and Linen

Flax: Linear Modules

Flax: Overview of Convolution and Sequence Models

Challenge: Flax

Introduction to Optax and Optimizers

Optax: Advanced Features

Challenge: Optax

Quiz: The JAX Ecosystem

Mini Project

Project: GAN Using the JAX ecosystem

Appendix

Installation Awesome JAX Libraries Awesome JAX-based Models Vector Calculus - II Common Tracer Errors Glossary

Introduction to JAX and Deep Learning/

...

/

Introduction to JAX Ecosystem

Introduction to JAX Ecosystem

This chapter provides an overview of the JAX Ecosystem.

We'll cover the following...

At this point in the course, JAX doesn’t seem much like a specialized neural network library, despite its wide suite of features. This is understandable given JAX, at ...