Repository: google/flax

Issue:
Improve `Module.init()` docstring
Right now the docstring of [`Module.init()`](https://github.com/google/flax/blob/df7fa7eed9965ab5117a50eee2813b8dc7c90426/flax/linen/module.py#L1351) is very short. It should explain much more and give more examples, including:

* Explain that users can pass both `{"params": rng}` as well as `rng` in case the only RNG used is the one for variable collection `params` (also link to the [Variables](https://flax.readthedocs.io/en/latest/api_reference/flax.linen.html#module-flax.core.variables) documentation).
* Give a few examples of how to use the various arguments.

We may have to change the docstring once #2428 is in as well.


Discussion:
Add me as a reviewer @zaxtax . We currently have to remind/teach users how to use `init` and `apply` in our notebooks/guides, while referencing the API docs for both (which don't go further i.t.o. explaining). This is a great Issue @marcvanzee .

Resolving patch:
diff --git a/flax/linen/module.py b/flax/linen/module.py
--- a/flax/linen/module.py
+++ b/flax/linen/module.py
@@ -1488,12 +1488,60 @@ def init(self,
            **kwargs) -> FrozenVariableDict:
     """Initializes a module method with variables and returns modified variables.
 
+    ``init`` takes as first argument either a single ``PRNGKey``, or a dictionary mapping variable collections names to their ``PRNGKeys``, and will call ``method`` (which is the module's ``__call__`` function by default) passing ``*args`` and ``**kwargs``, and returns
+    a dictionary of initialized variables.
+
+    Example::
+
+      >>> import flax.linen as nn
+      >>> import jax.numpy as jnp
+      >>> import jax
+      ...
+      >>> class Foo(nn.Module):
+      ...   @nn.compact
+      ...   def __call__(self, x, train):
+      ...     x = nn.Dense(16)(x)
+      ...     x = nn.BatchNorm(use_running_average=not train)(x)
+      ...     x = nn.relu(x)
+      ...     return nn.Dense(1)(x)
+      ...
+      >>> module = Foo()
+      >>> key = jax.random.PRNGKey(0)
+      >>> variables = module.init(key, jnp.empty((1, 7)), train=False)
+
+    If you pass a single ``PRNGKey``, Flax will use it to feed the ``'params'`` RNG stream.
+    If you want to use a different RNG stream or need to use multiple streams, you must pass a
+    dictionary mapping each RNG stream name to its corresponding ``PRNGKey`` to ``init``.
+
+    Example::
+
+      >>> class Foo(nn.Module):
+      ...   @nn.compact
+      ...   def __call__(self, x, train):
+      ...     x = nn.Dense(16)(x)
+      ...     x = nn.BatchNorm(use_running_average=not train)(x)
+      ...     x = nn.relu(x)
+      ...
+      ...     # Add gaussian noise
+      ...     noise_key = self.make_rng('noise')
+      ...     x = x + jax.random.normal(noise_key, x.shape)
+      ...
+      ...     return nn.Dense(1)(x)
+      ...
+      >>> module = Foo()
+      >>> rngs = {'params': jax.random.PRNGKey(0), 'noise': jax.random.PRNGKey(1)}
+      >>> variables = module.init(rngs, jnp.empty((1, 7)), train=False)
+
     Jitting `init` initializes a model lazily using only the shapes of the
     provided arguments, and avoids computing the forward pass with actual
     values. Example::
 
-      jit_init = jax.jit(SomeModule(...).init)
-      jit_init(rng, jnp.ones(input_shape, jnp.float32))
+      >>> module = nn.Dense(1)
+      >>> init_jit = jax.jit(module.init)
+      >>> variables = init_jit(jax.random.PRNGKey(0), jnp.empty((1, 7)))
+
+    ``init`` is a light wrapper over ``apply``, so other ``apply`` arguments like
+    ``method``, ``mutable``, and ``capture_intermediates`` are also available.
 
     Args:
       rngs: The rngs for the variable collections.


Test patch: