HomeCoursesIntroduction to JAX and Deep Learning

Intermediate

2h 30min

Updated 5 months ago

Introduction to JAX and Deep Learning
Save

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.
Join 2.7 million developers at
Overview
Content
Reviews
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.
JAX is a Python library designed for high-performance ML research. It is a powerful numerical computing library, just like Numpy...Show More

WHAT YOU'LL LEARN

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
Learn the basics of JAX

Show more

TAKEAWAY SKILLS

Random Variables

Neural Networks

Functional Programming

Deep learning basics

Content

1.

Introduction

2 Lessons

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

2.

JAX Programming Model

9 Lessons

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

6.

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
Every Educative resource is designed by our in-house team of ex-MAANG software engineers and PhD computer science educators — subject matter experts who’ve shipped production code at scale and taught the theory behind it. The goal is to get you hands-on with the skills you need to stay ahead in today's constantly evolving tech landscape. No videos, no fluff — just interactive, project-based learning with personalized feedback that adapts to your goals and experience.

Trusted by 2.7 million developers working at companies

Hands-on Learning Powered by AI

See how Educative uses AI to make your learning more immersive than ever before.

Instant Code Feedback

Evaluate and debug your code with the click of a button. Get real-time feedback on test cases, including time and space complexity of your solutions.

Adaptive Learning

Explain with AI

AI Code Mentor

Free Resources

FOR TEAMS

Interested in this course for your business or team?

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