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 00000000..bcb31c9c Binary files /dev/null and b/docs/source/_static/resources/jax_and_tpus_casml.pdf differ diff --git a/docs/source/_static/resources/jax_ecosystem_technical_report.pdf b/docs/source/_static/resources/jax_ecosystem_technical_report.pdf new file mode 100644 index 00000000..0c2b5e6f Binary files /dev/null and b/docs/source/_static/resources/jax_ecosystem_technical_report.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 00000000..57cdfac2 Binary files /dev/null and b/docs/source/_static/resources/megablocks_jax.pdf differ diff --git a/docs/source/_templates/navbar-top.html b/docs/source/_templates/navbar-top.html index 475d322f..128a7b62 100644 --- a/docs/source/_templates/navbar-top.html +++ b/docs/source/_templates/navbar-top.html @@ -7,6 +7,7 @@
  • diff --git a/docs/source/index.rst b/docs/source/index.rst index 7aac42ba..db032ef9 100644 --- a/docs/source/index.rst +++ b/docs/source/index.rst @@ -36,6 +36,15 @@ JAX AI Stack JAX_Vision_transformer JAX_time_series_classification +.. toctree:: + :hidden: + :caption: Learn more about the JAX AI Ecosystem + :maxdepth: 1 + + learn_ecosystem + jax_tpu_casml + megablocks + .. toctree:: :hidden: :caption: Developer resources diff --git a/docs/source/jax_tpu_casml.md b/docs/source/jax_tpu_casml.md new file mode 100644 index 00000000..3f564a82 --- /dev/null +++ b/docs/source/jax_tpu_casml.md @@ -0,0 +1,3 @@ +# 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). diff --git a/docs/source/learn_ecosystem.md b/docs/source/learn_ecosystem.md new file mode 100644 index 00000000..61ed5d5c --- /dev/null +++ b/docs/source/learn_ecosystem.md @@ -0,0 +1,3 @@ +# The JAX Ecosystem: A Modular, Scalable, and High-Performance ML Ecosystem (Technical Report) + +The JAX AI ecosystem extends the JAX numerical core with a collection of Google-backed composable libraries, evolving it into a robust, end-to-end, open-source platform for Machine Learning at extreme scales. Learn more about the JAX AI Ecosystem in the [*The JAX Ecosystem: A Modular, Scalable, and High-Performance ML Ecosystem* Technical Report (pdf)](_static/resources/jax_ecosystem_technical_report.pdf) document. diff --git a/docs/source/megablocks.md b/docs/source/megablocks.md new file mode 100644 index 00000000..78d9a710 --- /dev/null +++ b/docs/source/megablocks.md @@ -0,0 +1,5 @@ +# 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. diff --git a/docs/source/videos.md b/docs/source/videos.md index 18f75cf0..315b9dca 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.