MLP in JAX from Scratch

Implement initialization, forward passes, loss, autodiff, and pure functional SGD updates in JAX.

Easy1.8 hours21 steps

Overview

Implement initialization, forward passes, loss, autodiff, and pure functional SGD updates in JAX.

What you'll learn

JAX Autodiff Functional programming

FULL WALKTHROUGH

Learn every step on its own page

This project is no longer compressed into a few chapters. Open the dedicated learning workspace for a lesson-by-lesson explanation with concepts, MathJax mathematics, code, tests, mistakes, checkpoints, and persistent navigation.

Open 21-step walkthrough →

Build progress

0 / 21 steps0%

Move the tracker as you finish the original Deep-ML steps. Reaching 100% unlocks the completion action and certificate.

Architecture

Work through the system one dependable layer at a time. Each stage feeds the next and remains independently testable.

01PRNG keys→

02jit→

03grad

∑

Mathematics & visual explanation

Translate the core equations into code, validate intermediate tensors, and compare the implementation with a small numerical reference.

objective(θ) = data_term(θ) + λ · regularization(θ)

The exact objective evolves with each milestone. Keep a notebook of shapes, invariants, and numerical checks.