Merge pull request google#366 from fabianp:pjit_example

PiperOrigin-RevId: 500949450
michaelsdr · Jan 10, 2023 · 761fc96 · 761fc96
2 parents 4e33ee4 + bfd5ecf
commit 761fc96
Show file tree

Hide file tree

Showing 3 changed files with 39 additions and 0 deletions.
diff --git a/docs/notebooks/distributed/thumbnails/plot_custom_loop_pjit_example.png b/docs/notebooks/distributed/thumbnails/plot_custom_loop_pjit_example.png
diff --git a/docs/notebooks/distributed/thumbnails/plot_custom_loop_pmap_example.png b/docs/notebooks/distributed/thumbnails/plot_custom_loop_pmap_example.png
diff --git a/docs/notebooks/index.rst b/docs/notebooks/index.rst
@@ -45,6 +45,7 @@ Deep learning
 .. raw:: html
 
     <div class="sphx-glr-clear"></div>
+
 Implicit Differentiation
 ------------------------
 
@@ -80,3 +81,41 @@ Implicit Differentiation
 .. raw:: html
 
     </div>
+
+
+
+Distributed Optimization
+------------------------
+
+
+.. raw:: html
+
+    <div class="sphx-glr-thumbcontainer" tooltip="Distributed Optimization with JAXopt.">
+
+.. only:: html
+
+ .. figure:: /notebooks/distributed/thumbnails/plot_custom_loop_pjit_example.png
+     :alt: `jax.experimental.pjit` example using JAXopt.
+
+     :doc:`/notebooks/distributed/custom_loop_pjit_example`
+
+.. raw:: html
+
+    </div>
+
+
+
+.. raw:: html
+
+    <div class="sphx-glr-thumbcontainer" tooltip="pmap example using jaxopt">
+
+.. only:: html
+
+ .. figure:: /notebooks/distributed/thumbnails/plot_custom_loop_pmap_example.png
+     :alt: `jax.pmap` example using JAXopt.
+
+     :doc:`/notebooks/distributed/custom_loop_pmap_example`
+
+.. raw:: html
+
+    </div>