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

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

docs/source/stack_overview.md

@@ -0,0 +1,26 @@
+# Stack overview
+
+The JAX AI Stack is comprised of the following packages:
+
+* [JAX{material-regular}`open_in_new`](https://jax.dev): high-performance array
+  computing
+* [Flax
+  NNX{material-regular}`open_in_new`](https://flax.readthedocs.io/en/latest/):
+  object-oriented neural nets
+* [Optax{material-regular}`open_in_new`](https://optax.readthedocs.io/en/latest/index.html):
+  optimizers
+* [Orbax{material-regular}`open_in_new`](https://orbax.readthedocs.io/en/latest/):
+  checkpointing and model export
+* [Grain{material-regular}`open_in_new`](https://google-grain.readthedocs.io/en/latest/):
+  JAX-native data loading
+* [Chex{material-regular}`open_in_new`](https://chex.readthedocs.io/en/latest/):
+  JAX test utilities
+
+The `jax-ai-stack` metapackage installs compatible versions of all of these
+libraries, as well as shared compatible versions of shared dependencies.
+
+In addition, there is an optional `jax-ai-stack[tfds]` installation that
+includes [TensorFlow
+Datasets{material-regular}`open_in_new`](https://www.tensorflow.org/datasets),
+for those who wish to use TFDS for data loading. This includes a compatible
+version of TensorFlow as well.