DOC: make API doc titles more uniform

docs/jax.config.rst

@@ -1,22 +0,0 @@
-.. currentmodule:: jax
-
-JAX configuration
-=================
-
-.. autosummary::
-   :toctree: _autosummary
-
-   config
-   check_tracer_leaks
-   checking_leaks
-   debug_nans
-   debug_infs
-   default_device
-   default_matmul_precision
-   default_prng_impl
-   enable_checks
-   enable_custom_prng
-   enable_custom_vjp_by_custom_transpose
-   log_compiles
-   numpy_rank_promotion
-   transfer_guard

docs/jax.debug.rst

@@ -1,6 +1,6 @@
 
-jax.debug package
-=================
+``jax.debug`` module
+====================
 
 .. currentmodule:: jax.debug
 

docs/jax.distributed.rst

@@ -1,5 +1,5 @@
-jax.distributed module
-======================
+``jax.distributed`` module
+==========================
 
 .. currentmodule:: jax.distributed
 

docs/jax.dlpack.rst

@@ -1,5 +1,5 @@
-jax.dlpack module
-=================
+``jax.dlpack`` module
+=====================
 
 .. currentmodule:: jax.dlpack
 

docs/jax.example_libraries.optimizers.rst

@@ -1,5 +1,5 @@
-jax.example_libraries.optimizers module
-=======================================
+``jax.example_libraries.optimizers`` module
+===========================================
 
 .. automodule:: jax.example_libraries.optimizers
     :members:

docs/jax.example_libraries.rst

@@ -1,5 +1,5 @@
-jax.example_libraries package
-=============================
+``jax.example_libraries`` module
+================================
 
 JAX provides some small, experimental libraries for machine learning. These
 libraries are in part about providing tools and in part about serving as

docs/jax.example_libraries.stax.rst

@@ -1,5 +1,5 @@
-jax.example_libraries.stax module
-=================================
+``jax.example_libraries.stax`` module
+=====================================
 
 .. automodule:: jax.example_libraries.stax
     :members:

docs/jax.experimental.checkify.rst

@@ -1,5 +1,5 @@
-jax.experimental.checkify module
-=====================================
+``jax.experimental.checkify`` module
+====================================
 
 
 .. automodule:: jax.experimental.checkify

docs/jax.experimental.global_device_array.rst

@@ -1,5 +1,5 @@
-jax.experimental.global_device_array module
-===========================================
+``jax.experimental.global_device_array`` module
+===============================================
 
 .. automodule:: jax.experimental.global_device_array
 

docs/jax.experimental.host_callback.rst

@@ -1,5 +1,5 @@
-jax.experimental.host_callback module
-=====================================
+``jax.experimental.host_callback`` module
+=========================================
 
 
 .. automodule:: jax.experimental.host_callback

docs/jax.experimental.jet.rst

@@ -1,5 +1,5 @@
-jax.experimental.jet module
-===========================
+``jax.experimental.jet`` module
+===============================
 
 .. automodule:: jax.experimental.jet
 

docs/jax.experimental.maps.rst

@@ -1,5 +1,5 @@
-jax.experimental.maps module
-============================
+``jax.experimental.maps`` module
+================================
 
 .. automodule:: jax.experimental.maps
 

docs/jax.experimental.pjit.rst

@@ -1,5 +1,5 @@
-jax.experimental.pjit module
-============================
+``jax.experimental.pjit`` module
+================================
 
 .. automodule:: jax.experimental.pjit
 

docs/jax.experimental.rst

@@ -1,7 +1,7 @@
 .. currentmodule:: jax.experimental
 
-jax.experimental package
-========================
+``jax.experimental`` module
+===========================
 
 ``jax.experimental.optix`` has been moved into its own Python package
 (https://github.com/deepmind/optax).

docs/jax.experimental.sparse.rst

@@ -1,5 +1,5 @@
-jax.experimental.sparse module
-==============================
+``jax.experimental.sparse`` module
+==================================
 
 .. automodule:: jax.experimental.sparse
 

docs/jax.flatten_util.rst

@@ -1,5 +1,5 @@
-jax.flatten_util package
-========================
+``jax.flatten_util`` module
+===========================
 
 .. currentmodule:: jax.flatten_util
 

docs/jax.image.rst

@@ -1,5 +1,5 @@
-jax.image package
-=================
+``jax.image`` module
+====================
 
 .. currentmodule:: jax.image
 

docs/jax.lax.rst

@@ -1,5 +1,5 @@
-jax.lax package
-===============
+``jax.lax`` module
+==================
 
 .. automodule:: jax.lax
 

docs/jax.lib.rst

@@ -1,5 +1,5 @@
-jax.lib package
-===============
+``jax.lib`` module
+==================
 The `jax.lib` package is a set of internal tools and types for bridging between
 JAX's Python frontend and its XLA backend.
 

docs/jax.nn.initializers.rst

@@ -1,6 +1,5 @@
-
-jax.nn.initializers package
-===========================
+``jax.nn.initializers`` module
+==============================
 
 .. currentmodule:: jax.nn.initializers
 

docs/jax.nn.rst

@@ -1,5 +1,5 @@
 
-jax.nn package
+``jax.nn`` module
 =================
 
 .. currentmodule:: jax.nn
@@ -13,7 +13,7 @@ jax.nn package
 
 
 Activation functions
-------------------------
+--------------------
 
 .. autosummary::
   :toctree: _autosummary

docs/jax.numpy.rst

@@ -1,5 +1,5 @@
-jax.numpy package
-=================
+``jax.numpy`` module
+====================
 
 .. currentmodule:: jax.numpy
 

docs/jax.ops.rst

@@ -1,6 +1,5 @@
-
-jax.ops package
-===============
+``jax.ops`` module
+==================
 
 .. currentmodule:: jax.ops
 

docs/jax.profiler.rst

@@ -1,7 +1,7 @@
 .. currentmodule:: jax.profiler
 
-jax.profiler module
-===================
+``jax.profiler`` module
+=======================
 
 .. automodule:: jax.profiler
 

docs/jax.random.rst

@@ -1,5 +1,5 @@
-jax.random package
-==================
+``jax.random`` module
+=====================
 
 .. automodule:: jax.random
 

docs/jax.rst

@@ -7,32 +7,52 @@ Subpackages
 -----------
 
 .. toctree::
-    :maxdepth: 1
+   :maxdepth: 1
 
-    jax.numpy
-    jax.scipy
-    jax.sharding
-    jax.config
-    jax.debug
-    jax.dlpack
-    jax.distributed
-    jax.example_libraries
-    jax.experimental
-    jax.flatten_util
-    jax.image
-    jax.lax
-    jax.nn
-    jax.ops
-    jax.profiler
-    jax.random
-    jax.stages
-    jax.tree_util
+   jax.numpy
+   jax.scipy
+   jax.lax
+   jax.random
+   jax.sharding
+   jax.debug
+   jax.dlpack
+   jax.distributed
+   jax.flatten_util
+   jax.image
+   jax.nn
+   jax.ops
+   jax.profiler
+   jax.stages
+   jax.tree_util
+   jax.example_libraries
+   jax.experimental
 
 .. toctree::
    :hidden:
 
    jax.lib
 
+Configuration
+-------------
+
+.. autosummary::
+   :toctree: _autosummary
+
+   config
+   check_tracer_leaks
+   checking_leaks
+   debug_nans
+   debug_infs
+   default_device
+   default_matmul_precision
+   default_prng_impl
+   enable_checks
+   enable_custom_prng
+   enable_custom_vjp_by_custom_transpose
+   log_compiles
+   numpy_rank_promotion
+   transfer_guard
+
 .. _jax-jit:
 
 Just-in-time compilation (:code:`jit`)

docs/jax.scipy.rst

@@ -1,5 +1,5 @@
-jax.scipy package
-=================
+``jax.scipy`` module
+====================
 
 jax.scipy.fft
 -------------

docs/jax.sharding.rst

@@ -1,5 +1,5 @@
-jax.sharding package
-====================
+``jax.sharding`` module
+=======================
 
 .. automodule:: jax.sharding
 

docs/jax.stages.rst

@@ -1,5 +1,5 @@
-jax.stages package
-==================
+``jax.stages`` module
+=====================
 
 .. automodule:: jax.stages
 

docs/jax.tree_util.rst

@@ -1,5 +1,5 @@
-jax.tree_util package
-=====================
+``jax.tree_util`` module
+========================
 
 .. currentmodule:: jax.tree_util