From 12eebfe8d44fbd84e1cc0adc34c2a1ecf208df85 Mon Sep 17 00:00:00 2001 From: Roy Frostig Date: Tue, 13 Aug 2024 13:16:50 -0700 Subject: [PATCH] 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 | 22 +++------------------- docs/contributor_guide.rst | 22 ++++++++++++++++------ docs/distributed_data_loading.md | 2 +- docs/extensions.rst | 23 +++++++++++++++++++++++ docs/index.rst | 4 ++-- docs/jax.rst | 4 ++-- docs/jax_internal_api.rst | 4 ++-- docs/user_guides.rst | 22 ++++++++++++---------- 8 files changed, 61 insertions(+), 42 deletions(-) create mode 100644 docs/extensions.rst diff --git a/docs/advanced_guide.rst b/docs/advanced_guide.rst index cb987bd6e..5cf32f696 100644 --- a/docs/advanced_guide.rst +++ b/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 diff --git a/docs/contributor_guide.rst b/docs/contributor_guide.rst index b5ebd5057..55094fc88 100644 --- a/docs/contributor_guide.rst +++ b/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 diff --git a/docs/distributed_data_loading.md b/docs/distributed_data_loading.md index 70cbd26ba..be4d170ea 100644 --- a/docs/distributed_data_loading.md +++ b/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 diff --git a/docs/extensions.rst b/docs/extensions.rst new file mode 100644 index 000000000..92963b71f --- /dev/null +++ b/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 diff --git a/docs/index.rst b/docs/index.rst index 476de62bd..11d2807bf 100644 --- a/docs/index.rst +++ b/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 diff --git a/docs/jax.rst b/docs/jax.rst index 7be3e6301..79a46ff4d 100644 --- a/docs/jax.rst +++ b/docs/jax.rst @@ -1,7 +1,7 @@ .. currentmodule:: jax -Public API: jax package -======================= +Public API: ``jax`` package +=========================== Subpackages ----------- diff --git a/docs/jax_internal_api.rst b/docs/jax_internal_api.rst index fe65054d2..1ece596d8 100644 --- a/docs/jax_internal_api.rst +++ b/docs/jax_internal_api.rst @@ -1,5 +1,5 @@ -Internal APIs -============= +Internal API reference +====================== core ---- diff --git a/docs/user_guides.rst b/docs/user_guides.rst index 452600676..e917cf2fe 100644 --- a/docs/user_guides.rst +++ b/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