Merge pull request #2717 from marcvanzee:import-inits

PiperOrigin-RevId: 495556243

.gitignore

@@ -12,3 +12,4 @@ build/
 .pytype
 .vscode/*
 /.devcontainer
+docs/**/tmp

docs/api_reference/flax.linen.rst

@@ -149,27 +149,74 @@ Pooling
 Activation functions
 ------------------------
 
+.. automodule:: flax.linen.activation
+.. currentmodule:: flax.linen.activation
+
 .. autosummary::
   :toctree: _autosummary
 
+    PReLU
     celu
     elu
     gelu
     glu
+    hard_sigmoid
+    hard_silu
+    hard_swish
+    hard_tanh
+    leaky_relu
     log_sigmoid
     log_softmax
+    logsumexp
+    one_hot
     relu
+    relu6 as relu6,
+    selu
     sigmoid
+    silu
     soft_sign
     softmax
     softplus
+    standardize
     swish
-    PReLU
+    tanh
+
+
+Initializers
+------------------------
+
+.. automodule:: flax.linen.initializers
+.. currentmodule:: flax.linen.initializers
+
+.. autosummary::
+  :toctree: _autosummary
+
+    constant
+    delta_orthogonal
+    glorot_normal
+    glorot_uniform
+    he_normal
+    he_uniform
+    kaiming_normal
+    kaiming_uniform
+    lecun_normal
+    lecun_uniform
+    normal
+    ones
+    orthogonal
+    uniform
+    standardize
+    variance_scaling
+    xavier_normal
+    xavier_uniform
+    zeros
 
 
 Combinators
 ------------------------
 
+.. currentmodule:: flax.linen
+
 .. autosummary::
   :toctree: _autosummary
   :template: flax_module

flax/linen/__init__.py

@@ -41,6 +41,7 @@
   soft_sign as soft_sign,
   softmax as softmax,
   softplus as softplus,
+  standardize as standardize,
   swish as swish,
   tanh as tanh
 )

flax/linen/activation.py

@@ -44,6 +44,7 @@
 from jax.nn import soft_sign
 from jax.nn import softmax
 from jax.nn import softplus
+from jax.nn import standardize
 from jax.nn import swish
 import jax.numpy as jnp
 from jax.numpy import tanh