Merge pull request #33372 from jakevdp:jax-nn-docs

PiperOrigin-RevId: 833482568

Author: Google-ML-Automation

docs/jax.nn.initializers.rst

@@ -26,6 +26,8 @@ data type ``dtype``. Argument ``key`` is a PRNG key (e.g. from
     glorot_uniform
     he_normal
     he_uniform
+    kaiming_normal
+    kaiming_uniform
     lecun_normal
     lecun_uniform
     normal
@@ -34,4 +36,7 @@ data type ``dtype``. Argument ``key`` is a PRNG key (e.g. from
     truncated_normal
     uniform
     variance_scaling
+    xavier_normal
+    xavier_uniform
     zeros
+    Initializer

docs/jax.nn.rst

@@ -33,6 +33,7 @@ Activation functions
     hard_silu
     hard_swish
     hard_tanh
+    tanh
     elu
     celu
     selu

jax/_src/nn/initializers.py

@@ -49,6 +49,7 @@
 @export
 @typing.runtime_checkable
 class Initializer(Protocol):
+  """Protocol for initializers returned by :mod:`jax.nn.initializers` APIs."""
   def __call__(self,
                key: Array,
                shape: core.Shape,

tests/documentation_coverage_test.py

@@ -65,8 +65,6 @@ def jax_docs_dir() -> str:
   'jax.lax.linalg': [api for api in dir(jax.lax.linalg) if api.endswith('_p')],
   'jax.memory': ['Space'],
   'jax.monitoring': ['clear_event_listeners', 'record_event', 'record_event_duration_secs', 'record_event_time_span', 'record_scalar', 'register_event_duration_secs_listener', 'register_event_listener', 'register_event_time_span_listener', 'register_scalar_listener', 'unregister_event_duration_listener', 'unregister_event_listener', 'unregister_event_time_span_listener', 'unregister_scalar_listener'],
-  'jax.nn': ['tanh'],
-  'jax.nn.initializers': ['Initializer', 'kaiming_normal', 'kaiming_uniform', 'xavier_normal', 'xavier_uniform'],
   'jax.numpy': ['bfloat16', 'bool', 'e', 'euler_gamma', 'float4_e2m1fn', 'float8_e3m4', 'float8_e4m3', 'float8_e4m3b11fnuz', 'float8_e4m3fn', 'float8_e4m3fnuz', 'float8_e5m2', 'float8_e5m2fnuz', 'float8_e8m0fnu', 'inf', 'int2', 'int4', 'nan', 'newaxis', 'pi', 'uint2', 'uint4'],
   'jax.profiler': ['ProfileData', 'ProfileEvent', 'ProfileOptions', 'ProfilePlane', 'stop_server'],
   'jax.random': ['key_impl', 'random_gamma_p'],