[Pallas TPU] Add Pipelining and BlockSpecs documentation for Pallas TPU

docs/conf.py

@@ -122,6 +122,7 @@ exclude_patterns = [
     # These are kept in sync using the jupytext pre-commit hook.
     'notebooks/*.md',
     'pallas/quickstart.md',
+    'pallas/tpu/pipelining.md',
     'jep/9407-type-promotion.md',
     'jax-101/*.md',
     'autodidax.md',
@@ -211,6 +212,7 @@ nb_execution_excludepatterns = [
     'notebooks/autodiff_remat.*',
     # Requires accelerators
     'pallas/quickstart.*',
+    'pallas/tpu/pipelining.*',
 ]
 
 # -- Options for HTMLHelp output ---------------------------------------------

docs/pallas/index.rst

@@ -7,8 +7,14 @@ This section contains tutorials, guides and examples for using Pallas.
 
 .. toctree::
    :caption: Guides
-   :maxdepth: 1
+   :maxdepth: 2
 
    design
    quickstart
-   tpu
+
+.. toctree::
+   :caption: Platform Features
+   :maxdepth: 2
+
+   tpu/index
+

docs/pallas/tpu/index.rst

@@ -0,0 +1,10 @@
+Pallas TPU
+==========
+TPU specific documentation.
+
+.. toctree::
+   :caption: Guides
+   :maxdepth: 2
+
+   details
+   pipelining

docs/user_guides.rst

@@ -34,4 +34,8 @@ or deployed codebases.
    errors
    transfer_guard
 
+.. toctree::
+   :maxdepth: 1
+   :caption: Custom Operations
+
    pallas/index