cotalks.dev
Orgs
Login
Mastering the JAX AI Stack: From Flax NNX to Production
Channel:
Google for Developers
Videos (22)
1 — Leveraging the JAX AI Stack
2 — Optimizing Flax NNX Models with Optax (Part 1)
3 — Why JAX and Flax NNX?
4 — Introducing Flax NNX (Part 1)
5 — Introducing Flax NNX (Part 2)
6 — Debugging JAX & Flax NNX (Part 2)
7 — Introducing Flax NNX (Part 3)
8 — NumPy & JAX NumPy (Part 1)
9 — NumPy & JAX NumPy (Part 2)
10 — Debugging JAX & Flax NNX (Part 1)
11 — Getting started with MNIST
12 — Enhancing Reliability (Part 1)
13 — Scaling Up (Part 2)
14 — Efficient Data Loading
15 — Enhancing Reliability (Part 2)
16 — Scaling Up (Part 3)
17 — Scaling Up (Part 1)
18 — Optimizing Flax NNX Models with Optax (Part 2)
19 — Checkpointing Flax NNX Models with Orbax (Part 1)
20 — Serving JAX Models with vLLM & SGLang
21 — Checkpointing Flax NNX Models with Orbax (Part 2)
22 — JAX AI Stack: Summary & Conclusion
