Document jax.nn.initializers.

google · Mar 7, 2022 · d3d666d · d3d666d
1 parent 7061521
commit d3d666d
Show file tree

Hide file tree

Showing 3 changed files with 396 additions and 53 deletions.
diff --git a/docs/jax.nn.initializers.rst b/docs/jax.nn.initializers.rst
@@ -13,17 +13,25 @@ Initializers
 This module provides common neural network layer initializers,
 consistent with definitions used in Keras and Sonnet.
 
+An initializer is a function that takes three arguments:
+``(key, shape, dtype)`` and returns an array with dimensions ``shape`` and
+data type ``dtype``. Argument ``key`` is a :class:`jax.random.PRNGKey` random
+key used when generating random numbers to initialize the array.
+
 .. autosummary::
   :toctree: _autosummary
 
-    zeros
+    constant
+    delta_orthogonal
+    glorot_normal
+    glorot_uniform
+    he_normal
+    he_uniform
+    lecun_normal
+    lecun_uniform
+    normal
     ones
+    orthogonal
     uniform
-    normal
     variance_scaling
-    glorot_uniform
-    glorot_normal
-    lecun_uniform
-    lecun_normal
-    he_uniform
-    he_normal
+    zeros