Introduction to JAX and Deep Learning

Introduction to JAX and Deep Learning


53 Lessons

2h 30min

Certificate of Completion

Introduction to JAX and Deep Learning
1 Project
95 Playgrounds
14 Challenges
7 Quizzes
20 Illustrations

Takeaway Skills

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

Course Overview

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 sk...Show More


How You'll Learn

Hands-on Coding Environments

You don’t get better at swimming by watching others. Coding is no different. Practice as you learn with live code environments inside your browser.

2x Faster Learning — With No Setup

Videos are holding you back. Educative‘s interactive, text-based lessons accelerate learning — no setup, downloads, or alt-tabbing required.

AI-Powered Learning

Learn faster and smarter with adaptive AI tools embedded in every Educative course.

Progress You Can Show

Built-in assessments let you test your skills. Completion certificates let you show them off.


Interested in this course for your business or team?

Unlock this course (and 1,000+ more) for your entire org with DevPath