HomeLearning paths

What would you
like to build?

One shared foundation. Several ways to use it.

Start with JAX.

Arrays, functions and transformations. Build a checked regression experiment.

Open foundations path

Connect the lessons

Choose the system you want to build.

Follow training, recovery, evaluation and deployment in one project. Each guide identifies its runnable harness, available preparation and remaining target-runtime work.

01 / Runnable harness

Image: from pixels to a deployed classifier

Build an image classification system whose preprocessing, errors and runtime behavior you can inspect.

Explore the image guide

02 / Runnable harness

Text: from tokens to reproducible generation

Build a small causal language-model system with trustworthy masking, recovery and decoding.

Explore the text guide

03 / Runnable harness

Audio: from waveforms to an edge sound classifier

Build a sound-event system that preserves timing, handles silence and measures the full signal-processing path.

Explore the audio guide

04 / Runnable harness

Cross-modal: connect images and text

Build a retrieval system that checks alignment between modalities before optimizing similarity search.

Explore the cross-modal guide

Train models

Train neural networks

Train a CPU image classifier, preserve its complete state through interruption, inspect held-out errors and verify its exported inference. Language modeling and TPU work are optional branches.

Available lesson drafts · 53 lessons

First artifact: A CPU classifier with stable loss, replay checks and count-weighted evaluation.

19 available in the focus phases. Total above includes preparation.

Project available: A recoverable image classifier with verified exported inference

Open path

Reproducible JAX/Flax experiments workshop

Optional TPU and Google Cloud specialization

Optional language-model branch

Optional staged embedding-model pilot

Learn by interacting

Train a policy from actual environment interaction, verify its objective against an independent small-state reference, and compare multiple seeded agents on fixed held-out evaluation streams.

Available lesson drafts · 38 lessons

First artifact: A tested environment transition table and reproducible batched trajectories.

4 available in the focus phases. Total above includes preparation.

Project available: Train and audit a policy from interaction

Open path

Optional TPU and Google Cloud specialization

Learn TPUs and Google Cloud

Build and recover an experiment on CPU, then learn GCP provisioning, TPU generations, training and inference memory, precision policies and XProf profiling. Connect measured error and real target traces to model engineering.

Available lesson drafts · 87 lessons

First artifact: A verified TPU runtime receipt, a checkpoint-resumed TPU experiment log, and a memory/precision/XProf decision dossier.

35 available in the focus phases. Total above includes preparation.

Project available: Causal Transformer training-system capstone

Open path

Open the TPU & GCP practical guide

TPU generations, precision and profiling

Deploy models

Ship and adapt JAX workloads

Move trained computations into a measured inference service. Check framework/precision parity, package a non-root Docker image, verify readiness and connect evaluation, model identity and rollback evidence.

Available lesson drafts · 53 lessons

First artifact: An exported computation checked against its source on the same inputs.

8 available in the focus phases. Total above includes preparation.

Project available: Verify an inference artifact and capacity hypothesis

Open path

Optional TPU and Google Cloud specialization

Scientific work

Simulate and discover

Simulate physical dynamics, check numerical and derivative accuracy, and recover unknown parameters with held-out evidence. Build a reusable CPU inverse-problem audit before extending to adaptive solvers or accelerators.

Available lesson drafts · 33 lessons

First artifact: A cooling simulator with a convergence table and an independent sensitivity check.

4 available in the focus phases. Total above includes preparation.

Project available: Audit a differentiable scientific inverse problem

Open path

Model uncertainty

Express assumptions as probability models, verify a known posterior, then implement HMC and variational inference with explicit diagnostics. Use predictive checks to distinguish a sampler problem from a model that does not explain the data.

Available lesson drafts · 33 lessons

First artifact: An exact Bayesian regression posterior checked against an independent host calculation.

4 available in the focus phases. Total above includes preparation.

Project available: Audit a Bayesian regression model

Open path

Understand JAX itself

Trace a function into jaxpr, verify forward and reverse sensitivities, write exact custom derivative rules, and compile a small transformation whose limits you can explain.

Available lesson drafts · 33 lessons

First artifact: An annotated jaxpr that explains a computation and a changed input shape.

4 available in the focus phases. Total above includes preparation.

Project available: Audit derivatives and a tiny transformation

Open path

Systems work

Make it fast. Make it scale.

Profile a correct workload, preserve its global objective while partitioning arrays, and audit actual Pallas kernels before requesting target measurements. CPU semantics are available; accelerator speed and multi-host claims require their own receipts.

Available lesson drafts · 53 lessons

First artifact: A synchronized timing report and a checked array placement on logical CPU devices.

13 available in the focus phases. Total above includes preparation.

Project available: A checked kernel and a measurement-backed optimization argument

Open path

Optional TPU and Google Cloud specialization

Operate JAX workloads

Operate reproducible model workloads from tracked experiments to containerized inference. Practice failure recovery, MLOps release gates, LLMOps evaluation and tracing, and ModelOps ownership, approval and rollback.

Available lesson drafts · 65 lessons

First artifact: An interrupted CPU run whose next data batches, losses and full state match uninterrupted execution.

31 available in the focus phases. Total above includes preparation.

Project available: An observable, restartable local workload

Open path

Reproducible JAX/Flax experiments workshop

Optional TPU and Google Cloud specialization

Explore careers connected to these skills

See the work, milestones and portfolio goals behind each role.

ML & training engineer

Build models that train reliably.

  • Build and evaluate neural networks
  • Design data pipelines and resumable training
Read the career guide

Research engineer

Turn mathematical ideas into experiments.

  • Prototype differentiable algorithms
  • Investigate uncertainty and inference
Read the career guide

Scientific ML engineer

Simulate systems. Learn their hidden parameters.

  • Build and batch numerical simulations
  • Differentiate through dynamics and solvers
Read the career guide

RL engineer

Build agents that learn by interacting.

  • Build reproducible environments and rollouts
  • Implement and diagnose policy updates
Read the career guide

Accelerator & performance engineer

Find bottlenecks. Scale the right computation.

  • Diagnose compute, input, and communication bottlenecks
  • Partition workloads across devices
Read the career guide

Inference & deployment engineer

Make a trained model usable.

  • Adapt and export trained computations
  • Measure serving latency and throughput
Read the career guide

JAX / accelerator operations engineer

Keep real workloads observable and recoverable.

  • Manage accelerator jobs and runtime lifecycle
  • Monitor throughput, stalls, and checkpoint health
Read the career guide