DOC: organize content in advanced guide

docs/advanced_guide.rst

@@ -2,20 +2,41 @@
 
 Advanced Tutorials
 ==================
+This section contains examples and tutorials on more advanced topics, such as Multi Core computation, Custom operations, and more in depth applications
 
 .. toctree::
+   :caption: Examples
+   :maxdepth: 1
+
+   notebooks/neural_network_with_tfds_data
+   notebooks/Neural_Network_and_Data_Loading
+   notebooks/vmapped_log_probs
+
+.. toctree::
+   :caption: Parallel Computation
+   :maxdepth: 1
+
+   multi_process
+   notebooks/Distributed_arrays_and_automatic_parallelization
+   notebooks/xmap_tutorial
+
+.. toctree::
+   :caption: Automatic Differentiation
    :maxdepth: 1
 
    notebooks/autodiff_cookbook
-   multi_process
-   notebooks/Distributed_arrays_and_automatic_parallelization
-   notebooks/vmapped_log_probs
-   notebooks/neural_network_with_tfds_data
    notebooks/Custom_derivative_rules_for_Python_code
+
+.. toctree::
+   :caption: JAX Internals
+   :maxdepth: 1
+
    notebooks/How_JAX_primitives_work
    notebooks/Writing_custom_interpreters_in_Jax
-   notebooks/Neural_Network_and_Data_Loading
-   notebooks/xmap_tutorial
-   notebooks/convolutions
-   notebooks/external_callbacks
-   Custom_Operation_for_GPUs
+   Custom_Operation_for_GPUs
+
+.. toctree::
+   :caption: Deep Dives
+   :maxdepth: 1
+
+   notebooks/convolutions

docs/notebooks/convolutions.ipynb

@@ -6,7 +6,7 @@
     "id": "TVT_MVvc02AA"
    },
    "source": [
-    "# Convolutions in JAX\n",
+    "# Generalized Convolutions in JAX\n",
     "\n",
     "[![Open in Colab](https://colab.research.google.com/assets/colab-badge.svg)](https://colab.research.google.com/github/google/jax/blob/main/docs/notebooks/convolutions.ipynb)\n",
     "\n",

docs/notebooks/convolutions.md

@@ -14,7 +14,7 @@ kernelspec:
 
 +++ {"id": "TVT_MVvc02AA"}
 
-# Convolutions in JAX
+# Generalized Convolutions in JAX
 
 [![Open in Colab](https://colab.research.google.com/assets/colab-badge.svg)](https://colab.research.google.com/github/google/jax/blob/main/docs/notebooks/convolutions.ipynb)
 

docs/notebooks/vmapped_log_probs.ipynb

@@ -6,7 +6,7 @@
     "id": "6umP1IKf4Dg6"
    },
    "source": [
-    "# Autobatching log-densities example\n",
+    "# Autobatching for Bayesian Inference\n",
     "\n",
     "[![Open in Colab](https://colab.research.google.com/assets/colab-badge.svg)](https://colab.research.google.com/github/google/jax/blob/main/docs/notebooks/vmapped_log_probs.ipynb)\n",
     "\n",

docs/notebooks/vmapped_log_probs.md

@@ -14,7 +14,7 @@ kernelspec:
 
 +++ {"id": "6umP1IKf4Dg6"}
 
-# Autobatching log-densities example
+# Autobatching for Bayesian Inference
 
 [![Open in Colab](https://colab.research.google.com/assets/colab-badge.svg)](https://colab.research.google.com/github/google/jax/blob/main/docs/notebooks/vmapped_log_probs.ipynb)
 

docs/notes.rst

@@ -11,4 +11,5 @@ Notes
    concurrency
    gpu_memory_allocation
    rank_promotion_warning
-   jax_array_migration
+   jax_array_migration
+   type_promotion

docs/user_guides.rst

@@ -11,9 +11,9 @@ User Guides
    aot
    jaxpr
    pytrees
-   type_promotion
    errors
    debugging/index
    profiling
    device_memory_profiling
    transfer_guard
+   notebooks/external_callbacks