docs: reorganize sections

* Create "extension guides" section
* Sort developer notes into subsections
* Move examples from advanced section into user guides
* Reorder some listings, adjust some titles

docs/advanced_guide.rst

@@ -4,24 +4,16 @@ Advanced guides
 ===============
 
 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
+such as multi-core computation, automatic differentiation, and custom
+operations.
 
 .. toctree::
    :caption: Parallel computation
    :maxdepth: 1
 
-   multi_process
    notebooks/Distributed_arrays_and_automatic_parallelization
    notebooks/shard_map
+   multi_process
    distributed_data_loading
 
 .. toctree::
@@ -32,14 +24,6 @@ applications.
    notebooks/Custom_derivative_rules_for_Python_code
    notebooks/autodiff_remat
 
-.. toctree::
-   :caption: JAX internals
-   :maxdepth: 1
-
-   notebooks/How_JAX_primitives_work
-   notebooks/Writing_custom_interpreters_in_Jax
-   Custom_Operation_for_GPUs
-
 .. toctree::
    :caption: Deep dives
    :maxdepth: 1

docs/contributor_guide.rst

@@ -1,18 +1,28 @@
 .. _contributor-guide:
 
-Developer documentation
-=======================
+Developer notes
+===============
 
 JAX welcomes contributions from the community.
-See below for various install guides to get setup as a developer
-as well as developer-focused resources such as Jax Enhancement Proposals.
+These are guides to get set up as a developer, as well as
+developer-focused resources, such as JAX Enhancement Proposals.
+
+See also the :doc:`extension guides<../extensions>`, which document
+some of JAX's (extensible) internals.
+
 
 .. toctree::
    :maxdepth: 1
+   :caption: Contribution guides
 
    contributing
    developer
-   jax_internal_api
+   investigating_a_regression
+
+.. toctree::
+   :maxdepth: 1
+   :caption: Design and internals
+
    autodidax
    jep/index
-   investigating_a_regression
+   jax_internal_api

docs/distributed_data_loading.md

@@ -12,7 +12,7 @@ kernelspec:
   name: python3
 ---
 
-# Distributed data loading in a multi-host/multi-process environment
+# Distributed data loading in multi-host / multi-process environments
 
 <!--* freshness: { reviewed: '2024-05-16' } *-->
 

docs/extensions.rst

@@ -0,0 +1,23 @@
+.. _extensions:
+
+Extension guides
+================
+
+Guides for extending JAX's capabilities, and for building libraries
+that use or interface with JAX.
+
+.. toctree::
+   :caption: Extensible JAX internals
+   :maxdepth: 1
+
+   notebooks/How_JAX_primitives_work
+   jaxpr
+   notebooks/Writing_custom_interpreters_in_Jax
+   Custom_Operation_for_GPUs
+   jax.extend
+
+.. toctree::
+   :caption: Libraries and extensions
+   :maxdepth: 1
+
+   building_on_jax

docs/index.rst

@@ -48,7 +48,7 @@ For an end-to-end transformer library built on JAX, see MaxText_.
       :link-type: ref
       :class-card: user-guides
 
-    .. grid-item-card:: :material-regular:`laptop_chromebook;2em` Developer docs
+    .. grid-item-card:: :material-regular:`laptop_chromebook;2em` Developer notes
       :columns: 12 6 6 4
       :link: contributor-guide
       :link-type: ref
@@ -80,7 +80,7 @@ For an end-to-end transformer library built on JAX, see MaxText_.
    user_guides
    advanced_guide
    contributor_guide
-   building_on_jax
+   extensions
    notes
    jax
 

docs/jax.rst

@@ -1,7 +1,7 @@
 .. currentmodule:: jax
 
-Public API: jax package
-=======================
+Public API: ``jax`` package
+===========================
 
 Subpackages
 -----------

docs/jax_internal_api.rst

@@ -1,5 +1,5 @@
-Internal APIs
-=============
+Internal API reference
+======================
 
 core
 ----

docs/user_guides.rst

@@ -20,25 +20,27 @@ or deployed codebases.
 
 .. toctree::
    :maxdepth: 1
-   :caption: Development
+   :caption: Interfaces
 
-   jaxpr
-   notebooks/external_callbacks
-   type_promotion
    pytrees
-
-.. toctree::
-   :maxdepth: 1
-   :caption: Run time
-
+   errors
    aot
    export/index
-   errors
+   type_promotion
    transfer_guard
 
 .. toctree::
    :maxdepth: 1
    :caption: Custom operations
 
+   notebooks/external_callbacks
    pallas/index
    ffi
+
+.. toctree::
+   :caption: Example applications
+   :maxdepth: 1
+
+   notebooks/neural_network_with_tfds_data
+   notebooks/Neural_Network_and_Data_Loading
+   notebooks/vmapped_log_probs