HomeCoursesIntroduction to JAX and Deep Learning

AI-powered learning

Save

Introduction to JAX and Deep Learning

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.

4.8

53 Lessons

2h 30min

Join 2.9 million developers at

LEARNING OBJECTIVES

Learn the basics of JAX
Learn how to apply Autograd
Use auto vectorization for batching
Use Haiku and Flax for implementing neural networks
Cover Optax and overview of common optimization algorithms in deep learning
Use Chex for testing JAX programs
Learn the basics of applied linear algebra
Learn random variables theory and probability distributions
Learn pseudo-random number generation
Cover the basics of optimal transport

Learning Roadmap

53 Lessons1 Project7 Quizzes14 Challenges

Introduction

Get familiar with JAX, a powerful library for deep learning and numerical computing.

Course Introduction

JAX Overview

JAX Programming Model

Walk through JAX's programming model, including pure functions, JIT, jaxpr, and autodiff.

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

15 Lessons

Explore the fundamental concepts of vectors, matrices, multivariate calculus, and convolutions in deep learning.

Random Variables and Distributions

7 Lessons

Grasp the fundamentals of random variables, distributions, PRNGs, and divergence measures in JAX.

JAX Ecosystem

14 Lessons

Take a closer look at the tools and libraries within the JAX ecosystem for deep learning.

Project: GAN Using the JAX ecosystem

Project

Appendix

6 Lessons

Focus on installation steps, notable JAX libraries, models, vector calculus, common errors, and key terms.

Certificate of Completion

Showcase your accomplishment by sharing your certificate of completion.

Developed by MAANG Engineers

ABOUT THIS COURSE

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.

ABOUT THE AUTHOR

Khayyam Hashmi

Computer scientist and Generative AI and Machine Learning specialist. VP of Technical Content @ educative.io.

Learn more about Khayyam

Trusted by 2.9 million developers working at companies

These are high-quality courses. Trust me the price is worth it for the content quality. Educative came at the right time in my career. I'm understanding topics better than with any book or online video tutorial I've done. Truly made for developers. Thanks

Anthony Walker

@_webarchitect_

Just finished my first full #ML course: Machine learning for Software Engineers from Educative, Inc. ... Highly recommend!

Evan Dunbar

ML Engineer

You guys are the gold standard of crash-courses... Narrow enough that it doesn't need years of study or a full blown book to get the gist, but broad enough that an afternoon of Googling doesn't cut it.

Software Developer

Carlos Matias La Borde

I spend my days and nights on Educative. It is indispensable. It is such a unique and reader-friendly site

Souvik Kundu

Front-end Developer

Your courses are simply awesome, the depth they go into and the breadth of coverage is so good that I don't have to refer to 10 different websites looking for interview topics and content.

Vinay Krishnaiah

Software Developer