diff --git a/docs/jax.extend.rst b/docs/jax.extend.rst
index 1882d9a85..a846cf93a 100644
--- a/docs/jax.extend.rst
+++ b/docs/jax.extend.rst
@@ -2,3 +2,20 @@
 =====================
 
 .. automodule:: jax.extend
+
+``jax.extend.random``
+---------------------
+
+.. automodule:: jax.extend.random
+
+.. autosummary::
+  :toctree: _autosummary
+
+  PRNGImpl
+  seed_with_impl
+  threefry2x32_p
+  threefry_2x32
+  threefry_prng_impl
+  rbg_prng_impl
+  unsafe_rbg_prng_impl
+  wrap_key_data
diff --git a/docs/jax.random.rst b/docs/jax.random.rst
index 24793f97a..b22d5709c 100644
--- a/docs/jax.random.rst
+++ b/docs/jax.random.rst
@@ -16,6 +16,8 @@ List of Available Functions
   :toctree: _autosummary
 
     PRNGKey
+    key
+    key_data
     ball
     bernoulli
     beta