Class Central is learner-supported. When you buy through links on our site, we may earn an affiliate commission.

YouTube

Machine Learning with JAX - From Hero to HeroPro+

Aleksa Gordić - The AI Epiphany via YouTube

Overview

Dive into advanced machine learning concepts with JAX in this comprehensive tutorial video. Learn to convert stateful models to stateless, master PyTrees, and train a multilayer perceptron using pure JAX. Explore custom PyTrees, parallelism with TPUs, and inter-device communication. Discover techniques for training models across multiple machines, implementing per-example gradients, and even tackle meta-learning with a 3-line MAML implementation. Gain practical insights into JAX's powerful features for building and optimizing complex machine learning models.

Syllabus

My get started with JAX repo
Stateful to stateless conversion
PyTrees in depth
Training an MLP in pure JAX
Custom PyTrees
Parallelism in JAX TPUs example
Communication between devices
value_and_grad and has_aux
Training an ML model on multiple machines
stop grad, per example grads
Implementing MAML in 3 lines
Outro

Taught by

Aleksa Gordić - The AI Epiphany

Reviews

Start your review of Machine Learning with JAX - From Hero to HeroPro+

Never Stop Learning.

Get personalized course recommendations, track subjects and courses with reminders, and more.

Someone learning on their laptop while sitting on the floor.