Learning resources

Intro to JAX & TPUs

Learn more about TPUs, how they work, and how to use JAX and Pallas to use TPUs effectively in this video.

Megablocks

Megablocks is a library for building large-scale machine learning models with JAX.

Check out this video for an overview of Megablocks and its capabilities, including dealing with sparsity and ragged dot operations.

More learning resources

In https://goo.gle/learning-jax, you can find a curated list of learning resources about JAX and the JAX AI Stack, including slide decks, coding exercises and quick references.