Add/update JAX Advanced Tutorials docs, ToC structure

jax-ml · Sep 11, 2024 · 01e315b · 01e315b
1 parent c708b7c
commit 01e315b
Show file tree

Hide file tree

Showing 20 changed files with 21 additions and 4,429 deletions.
diff --git a/docs/_tutorials/index.rst b/docs/_tutorials/index.rst
@@ -38,10 +38,7 @@ JAX 201
    :maxdepth: 1
 
    parallelism
-   advanced-autodiff
-   gradient-checkpointing
    advanced-debugging
-   external-callbacks
    profiling-and-performance
 
 JAX 301
@@ -50,6 +47,4 @@ JAX 301
 .. toctree::
    :maxdepth: 1
 
-   jax-primitives
-   jaxpr
    advanced-compilation
diff --git a/docs/_tutorials/advanced-autodiff.md → docs/advanced-autodiff.md b/docs/_tutorials/advanced-autodiff.md → docs/advanced-autodiff.md
diff --git a/docs/conf.py b/docs/conf.py
@@ -85,7 +85,6 @@ def _do_not_evaluate_in_jax(
 ]
 
 intersphinx_mapping = {
-    'array_api': ('https://data-apis.org/array-api/2023.12/', None),
     'python': ('https://docs.python.org/3/', None),
     'numpy': ('https://numpy.org/doc/stable/', None),
     'scipy': ('https://docs.scipy.org/doc/scipy/reference/', None),
@@ -358,4 +357,6 @@ def linkcode_resolve(domain, info):
     'jax-101/07-state.md': 'stateful-computations.md',
     'jax-101/08-pjit.rst': 'sharded-computation.md',
     'jax-101/index.rst': 'tutorials.rst',
+    'notebooks/external_callbacks.md': 'external-callbacks.md',
+    'notebooks/How_JAX_primitives_work.md': 'jax-primitives.md',
 }
diff --git a/docs/extensions.rst b/docs/extensions.rst
@@ -10,8 +10,6 @@ that use or interface with JAX.
    :caption: Extensible JAX internals
    :maxdepth: 1
 
-   notebooks/How_JAX_primitives_work
-   jaxpr
    notebooks/Writing_custom_interpreters_in_Jax
    Custom_Operation_for_GPUs
    jax.extend

diff --git a/docs/_tutorials/external-callbacks.md → docs/external-callbacks.md b/docs/_tutorials/external-callbacks.md → docs/external-callbacks.md
diff --git a/docs/ffi.ipynb b/docs/ffi.ipynb
@@ -4,8 +4,6 @@
    "cell_type": "markdown",
    "metadata": {},
    "source": [
-    "(ffi-tutorial)=\n",
-    "\n",
     "# Foreign function interface (FFI)\n",
     "\n",
     "_This tutorial requires JAX v0.4.31 or newer._\n",
@@ -364,7 +362,7 @@
    "cell_type": "markdown",
    "metadata": {},
    "source": [
-    "We can inspect the [jaxpr](understanding-jaxprs) of the {func}`~jax.vmap` of `rms_norm` to confirm that it isn't being rewritten using {func}`~jax.lax.scan`:"
+    "We can inspect the [jaxpr](jax-internals-jaxpr) of the {func}`~jax.vmap` of `rms_norm` to confirm that it isn't being rewritten using {func}`~jax.lax.scan`:"
    ]
   },
   {

diff --git a/docs/ffi.md b/docs/ffi.md
@@ -311,7 +311,7 @@ Our implementation of `rms_norm` has the appropriate semantics, and it supports
 np.testing.assert_allclose(jax.vmap(rms_norm)(x), jax.vmap(rms_norm_ref)(x), rtol=1e-5)
 ```
 
-We can inspect the [jaxpr](understanding-jaxprs) of the {func}`~jax.vmap` of `rms_norm` to confirm that it isn't being rewritten using {func}`~jax.lax.scan`:
+We can inspect the [jaxpr](jax-internals-jaxpr) of the {func}`~jax.vmap` of `rms_norm` to confirm that it isn't being rewritten using {func}`~jax.lax.scan`:
 
 ```{code-cell} ipython3
 jax.make_jaxpr(jax.vmap(rms_norm))(x)

diff --git a/docs/glossary.rst b/docs/glossary.rst
@@ -30,7 +30,7 @@ Glossary of terms
     jaxpr
       Short for *JAX expression*, a jaxpr is an intermediate representation of a computation that
       is generated by JAX, and is forwarded to :term:`XLA` for compilation and execution.
-      See :ref:`understanding-jaxprs` for more discussion and examples. 
+      See :ref:`jax-internals-jaxpr` for more discussion and examples. 
 
     JIT
       Short for *Just In Time* compilation, JIT in JAX generally refers to the compilation of

diff --git a/docs/_tutorials/gradient-checkpointing.md → docs/gradient-checkpointing.md b/docs/_tutorials/gradient-checkpointing.md → docs/gradient-checkpointing.md
diff --git a/docs/index.rst b/docs/index.rst
@@ -121,20 +121,21 @@ maintains an up-to-date list.
 
    installation
    quickstart
-   notebooks/Common_Gotchas_in_JAX
-   faq
 
 .. toctree::
    :hidden:
    :maxdepth: 1
 
    tutorials
 
+   notebooks/Common_Gotchas_in_JAX
+
+   faq
 
 .. toctree::
    :hidden:
    :maxdepth: 2
-   :caption: Resources
+   :caption: More guides/resources
 
    user_guides
    advanced_guide

diff --git a/docs/_tutorials/jax-primitives.md → docs/jax-primitives.md b/docs/_tutorials/jax-primitives.md → docs/jax-primitives.md
diff --git a/docs/_tutorials/jaxpr.md → docs/jaxpr.md b/docs/_tutorials/jaxpr.md → docs/jaxpr.md