diff --git a/docs/source/_static/resources/jax_and_tpus_casml.pdf b/docs/source/_static/resources/jax_and_tpus_casml.pdf new file mode 100644 index 0000000..bcb31c9 Binary files /dev/null and b/docs/source/_static/resources/jax_and_tpus_casml.pdf differ diff --git a/docs/source/_static/resources/megablocks_jax.pdf b/docs/source/_static/resources/megablocks_jax.pdf new file mode 100644 index 0000000..57cdfac Binary files /dev/null and b/docs/source/_static/resources/megablocks_jax.pdf differ diff --git a/docs/source/index.rst b/docs/source/index.rst index 7aac42b..13683e4 100644 --- a/docs/source/index.rst +++ b/docs/source/index.rst @@ -11,7 +11,6 @@ JAX AI Stack install getting_started - more_learning .. toctree:: :hidden: @@ -52,3 +51,4 @@ JAX AI Stack news events videos + more_learning diff --git a/docs/source/more_learning.md b/docs/source/more_learning.md index ff7f2e2..025527d 100644 --- a/docs/source/more_learning.md +++ b/docs/source/more_learning.md @@ -1,3 +1,15 @@ -# More Learning Resources +# 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 [Intro to JAX & TPUs slide deck](_static/resources/jax_and_tpus_casml.pdf) and [video](https://rsvp.withgoogle.com/events/intro-to-jax-and-google-tpu-for-casml-attendees/sessions/inro-to-jax-and-google-tpu). + +## Megablocks + +[Megablocks](https://arxiv.org/pdf/2211.15841) is a library for building large-scale machine learning models with JAX. + +Check out [this presentation](_static/resources/megablocks_jax.pdf) and watch [this video](https://www.youtube.com/watch?v=tWkMj6lUp1c) 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. diff --git a/docs/source/videos.md b/docs/source/videos.md index 18f75cf..315b9dc 100644 --- a/docs/source/videos.md +++ b/docs/source/videos.md @@ -3,3 +3,5 @@ A collection of videos related to the JAX AI Stack, including talks, tutorials, and presentations. - [JAX DevLab 2025 Playlist](https://www.youtube.com/playlist?list=PLlFotmaRrOztzhilpCT4vBsa0R9hz5ADD) +- Learn more about TPUs, how they work, and how to use JAX and Pallas to use TPUs effectively in this [video](https://rsvp.withgoogle.com/events/intro-to-jax-and-google-tpu-for-casml-attendees/sessions/inro-to-jax-and-google-tpu). +- Watch [this video](https://www.youtube.com/watch?v=tWkMj6lUp1c) for an overview of [Megablocks](https://arxiv.org/pdf/2211.15841), a library for building large-scale machine learning models with JAX.