[WIP] Add "The stack" section to left nav

The new section contains an overview of the libraries in the stack, as
well as a page for each library with a brief description and outbound
links for more information.

Author: skye

docs/requirements.txt

@@ -5,6 +5,7 @@ myst-nb
 myst-parser[linkify]
 sphinx-book-theme
 sphinx-copybutton
+sphinx-design
 
 # Packages required for notebook execution
 matplotlib

docs/source/conf.py

@@ -17,6 +17,7 @@
 extensions = [
     'myst_nb',
     'sphinx_copybutton',
+    'sphinx_design',
 ]
 
 templates_path = ['_templates']

docs/source/index.rst

@@ -12,6 +12,20 @@ JAX AI Stack
    install
    getting_started
 
+.. toctree::
+   :hidden:
+   :caption: The stack
+   :maxdepth: 1
+
+   stack_overview
+   stack_jax
+   stack_flax
+   stack_optax
+   stack_orbax_checkpoint
+   stack_orbax_export
+   stack_grain
+   stack_chex
+
 .. toctree::
    :hidden:
    :caption: Tutorials

docs/source/stack_chex.md

@@ -0,0 +1 @@
+# Chex: test utilities

docs/source/stack_flax.md

@@ -0,0 +1,24 @@
+# Flax NNX: neural nets
+
+Flax NNX provides **neural net functionality** on top of JAX, such as a module
+abstraction and pre-defined layers, via a **Pythonic object-oriented API**. NNX
+allows you to write stateful model code that can still take advantage of JAX's
+function transforms and other features.
+
+NNX has native integration with [Optax](stack_optax).
+
+Main Flax NNX site:
+**[flax.readthedocs.io{material-regular}`open_in_new`](https://flax.readthedocs.io/)**
+
+**If you'd like to learn more about NNX** beyond what's covered in the
+[](getting_started) guide, we recommend starting with **[Flax
+basics{material-regular}`open_in_new`](https://flax.readthedocs.io/en/latest/nnx_basics.html)**.
+
+The Flax NNX docs cover many other useful topics including:
+
+* [Function
+  transforms{material-regular}`open_in_new`](https://flax.readthedocs.io/en/latest/guides/transforms.html)
+* [Parallelism{material-regular}`open_in_new`](https://flax.readthedocs.io/en/latest/guides/flax_gspmd.html)
+* [Performance
+  considerations{material-regular}`open_in_new`](https://flax.readthedocs.io/en/latest/guides/performance.html)
+* And much more!

docs/source/stack_grain.md

@@ -0,0 +1 @@
+# Grain: data loading

docs/source/stack_jax.md

@@ -0,0 +1,18 @@
+# JAX: array computing
+
+JAX is the foundation of the JAX AI Stack! It provides **high-performance array
+computing** functionality over accelerators via a simple **NumPy-like API and
+function transformations**.
+
+Main JAX site: **[jax.dev{material-regular}`open_in_new`](https://jax.dev)**
+
+**If you'd like to learn more about JAX** beyond what's covered in the
+[](getting_started) guide, we recommend starting with the **[JAX
+tutorials{material-regular}`open_in_new`](https://docs.jax.dev/en/latest/tutorials.html)**.
+
+The JAX docs cover many other useful topics including:
+
+* [Performance profiling{material-regular}`open_in_new`](https://docs.jax.dev/en/latest/profiling.html)
+* [Multi-host JAX programs{material-regular}`open_in_new`](https://docs.jax.dev/en/latest/multi_process.html)
+* [Custom GPU + TPU kernels with Pallas{material-regular}`open_in_new`](https://docs.jax.dev/en/latest/pallas/index.html)
+* And much more!

docs/source/stack_optax.md

@@ -0,0 +1,12 @@
+# Optax: optimizers
+
+Optax provides **gradient processing and optimization** functionality on top of
+JAX, including optimizers and losses.
+
+Main Optax site:
+**[optax.readthedocs.io{material-regular}`open_in_new`](https://optax.readthedocs.io/en/latest/index.html)**
+
+**If you'd like to learn more about Optax** beyond what's covered in the
+[](getting_started) guide, we recommend starting with the **[Optax getting
+started{material-regular}`open_in_new`](https://optax.readthedocs.io/en/latest/getting_started.html)**
+guide.

docs/source/stack_orbax_checkpoint.md

@@ -0,0 +1 @@
+# Orbax: checkpointing

docs/source/stack_orbax_export.md

@@ -0,0 +1 @@
+# Orbax: model export