Repository: google/jax

Issue:
JAX Random Docs Page is Blank
Seems the page for Jax.random is not generating correctly...?

https://jax.readthedocs.io/en/latest/jax.random.html


Discussion:
Thanks for reporting. I'll take a look - I suspect it may be because jax/random.py does not have a module-level docstring.
Hmm... looks like it's due to more than just the docstring. #4964 put the random functions into a private module, so sphinx sees this and skips them in autodoc. The solution, I think, is to explicitly list all functions to be documented, as we do in `jax.numpy.rst`.

Resolving patch:
diff --git a/jax/_src/random.py b/jax/_src/random.py
--- a/jax/_src/random.py
+++ b/jax/_src/random.py
@@ -12,32 +12,6 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 
-"""JAX pseudo-random number generators (PRNGs).
-
-Example usage:
-
->>> rng = jax.random.PRNGKey(seed)
->>> for i in range(num_steps):
-...   rng, rng_input = jax.random.split(rng)
-...   params = compiled_update(rng_input, params, next(batches))
-
-Context:
-
-Among other requirements, the JAX PRNG aims to:
-(a) ensure reproducibility,
-(b) parallelize well, both in terms of vectorization (generating array values)
-and multi-replica, multi-core computation. In particular it should not use
-sequencing constraints between random function calls.
-
-The approach is based on:
-1. "Parallel random numbers: as easy as 1, 2, 3" (Salmon et al. 2011)
-2. "Splittable pseudorandom number generators using cryptographic hashing"
-(Claessen et al. 2013)
-
-See also https://github.com/google/jax/blob/master/design_notes/prng.md
-for the design and its motivation.
-"""
-
 
 from functools import partial
 from typing import Optional, Sequence, Union
diff --git a/jax/random.py b/jax/random.py
--- a/jax/random.py
+++ b/jax/random.py
@@ -12,6 +12,68 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 
+"""Utilities for pseudo-random number generation.
+
+The ``jax.random`` package provides a number of routines for deterministic
+generation of sequences of pseudorandom numbers.
+
+Basic usage
+-----------
+
+>>> key = jax.random.PRNGKey(seed)
+>>> for i in range(num_steps):
+...   key, subkey = jax.random.split(key)
+...   params = compiled_update(subkey, params, next(batches))
+
+PRNG Keys
+---------
+Unlike the *stateful* pseudorandom number generators (PRNGs) that users of NumPy and
+SciPy may be accustomed to, JAX random functions all require an explicit PRNG state to
+be passed as a first argument.
+The random state is described by two unsigned 32-bit integers that we call a **key**,
+usually generated by the :py:func:`jax.random.PRNGKey` function::
+
+    >>> from jax import random
+    >>> key = random.PRNGKey(0)
+    >>> key
+    DeviceArray([0, 0], dtype=uint32)
+
+This key can then be used in any of JAX's random number generation routines::
+
+    >>> random.uniform(key)
+    DeviceArray(0.41845703, dtype=float32)
+
+Note that using a key does not modify it, so reusing the same key will lead to the same result::
+
+    >>> random.uniform(key)
+    DeviceArray(0.41845703, dtype=float32)
+
+If you need a new random number, you can use :meth:`jax.random.split` to generate new subkeys::
+
+    >>> key, subkey = random.split(key)
+    >>> random.uniform(subkey)
+    DeviceArray(0.10536897, dtype=float32)
+
+Design and Context
+------------------
+
+Among other requirements, the JAX PRNG aims to:
+
+(a) ensure reproducibility,
+(b) parallelize well, both in terms of vectorization (generating array values)
+    and multi-replica, multi-core computation. In particular it should not use
+    sequencing constraints between random function calls.
+
+The approach is based on:
+
+1. "Parallel random numbers: as easy as 1, 2, 3" (Salmon et al. 2011)
+2. "Splittable pseudorandom number generators using cryptographic hashing"
+   (Claessen et al. 2013)
+
+See also https://github.com/google/jax/blob/master/design_notes/prng.md
+for the design and its motivation.
+"""
+
 # flake8: noqa: F401
 
 from jax._src.random import (


Test patch: