Solving y=mx+b with Jax on a TPU Pod slice - Mat Kelcey

matpalm.comsite

A four part YouTube tutorial series with Colab notebooks that starts with Jax fundamentals and moves up to training with a data parallel approach on a v3-32 TPU Pod slice.

from

JAX

added

2026-10-10

likes

0

similar

1. ▲

   JAX github.com

   Composable transformations of Python+NumPy programs: differentiate, vectorize, JIT to GPU/TPU, and more.

   tool topic programming link ↗
   
2. ▲

   Jax in Action manning.com

   A hands-on guide to using JAX for deep learning and other mathematically-intensive applications.

   site topic ai link ↗
   site
3. ▲

   SymJAX github.com

   Symbolic CPU/GPU/TPU programming.

   tool topic programming link ↗
   
4. ▲

   JAX: Accelerated Machine Learning Research | SciPy 2020 | VanderPlas youtube.com

   JAX's core design, how it's powering new research, and how you can start using it.

   video · 24 min · Enthought topic aitopic design link ↗
   
5. ▲

   GPJax github.com

   Gaussian processes in JAX.

   tool link ↗
   
6. ▲

   JAX Toolbox github.com

   Nightly CI and optimized examples for JAX on NVIDIA GPUs using libraries such as T5x, Paxml, and Transformer Engine.

   tool topic ai link ↗
   

JAX › Videos: “A four part YouTube tutorial series with Colab notebooks that starts with Jax fundamentals and moves up to training with a data parallel approach on a v3-32 TPU Pod slice.”