Repository: google/flax

Issue:
Full Module RNG Scope
Converting this to an issue since I'm repeatedly stumbling into the current non-streamlined implementation for module random-key management (see discussion below; it doesn't present a solution yet, but it writes down some ideas) .

The biggest issue that the current design has is that it prevents me from composing generic modules without having to awkardly pass a scope of RNGs (and likely make a mistake). 

This snippet shows a example use-case for sampling and applying random functions,
```python
from typing import Generic, TypeVar

import jax
import jax.numpy as jnp

from flax import struct
from flax import linen as nn


M = TypeVar('M', bound=nn.Module)

@struct.dataclass
class Composite(Generic[M]):
    module: M

    def do(self, key: jax.Array, ins: jax.Array):
        # Sample random module parameters, and apply the function.
        
        key_init, keys_sample = jax.random.split(key)
        
        # Problem: Throws an error if an arbitrary key in `M` is not covered.
        params = self.module.init(key, ins)

        # Problem: I do not know the RNG scope of self.module
        outs = self.module.apply(params, ins)
        
        return outs


c = Composite(nn.Dense(5))
c.do(jax.random.key(0), jnp.ones(5))
> Array([-0.7242578 ,  0.50101703,  0.44527686,  0.31400478, -0.02983775],      dtype=float32)
```

I'd personally suggest that flax should silently split up `key` inside the `init` and `apply` for me. Another solution is that flax returns all observed rngs in the Module scope.


### Discussed in https://github.com/google/flax/discussions/3375

<div type='discussions-op-text'>

<sup>Originally posted by **joeryjoery** September 28, 2023</sup>
Hi, if I implement a Module that makes use of `make_rng`, e.g., Dropout, I now have to manually specify the `rng` when I use the Module's call function. Though I completely agree that the RNG streams should be split into distinct paths like it is now, it is quite annoying that I have to bookkeep this at every call to `Module.apply`.

```python
import jax
import flax.linen as nn


rng = jax.random.PRNGKey(0)


class A(nn.Module):

    def __call__(self):
        x = self.make_rng('x')
        y = self.make_rng('y')

        return jax.random.uniform(x, ()) + jax.random.uniform(y, ())

myA = A()

# Currently
keys = jax.random.split(rng, 3)
scope = {'params', 'x', 'y'}  # I don't want to do this manually

# Desired functionality (or something like this)
# scope = myA.get_rng_scope()

rngs = dict(zip(scope, keys))

v = myA.init(rngs)
_ = myA.apply(v, rngs=rngs)
```
Is there already something like this implemented? Or could this be done in the future?

I'd also argue that flax should allow me to just pass in one global RNG key to either `init` or `apply`.
```python
# Even better (just use one key)
_ = myA.apply(v, rngs=jax.random.PRNGKey(0))
```
This functionality would be great to have with the *option* to pass in the RNGS explicitly. This has no consequences for reproducibility since the rngs need to be split somewhere anyway. 

At the moment the RNG design forces me to write work-arounds in my code which requires me to pass around a namespace that tracks all registered scopes. But this is quite error-prone and ideally this would be handled internally.</div>


Discussion:
What if you explicitly return the scope?
```
class MLP(nn.Module):
  @nn.compact
  def __call__(self, x, train):
    x = nn.Dense(2)(x)
    x = nn.Dense(3)(x)
    x = nn.Dropout(0.1)(x, deterministic=not train)
    return x
  def get_scope_rngs(self, x, train):
    x = self.__call__(x, train)
    return x, self.scope.rngs

model = MLP()
x = jnp.ones((1, 2))
(out, scope_rngs), v = model.init_with_output({'params': jax.random.key(0), 'dropout': jax.random.key(1)}, x, train=False, method='get_scope_rngs')
key_dict = {k: jax.random.key(i) for i, k in enumerate(scope_rngs.keys())}
out = model.apply(v, x, train=True, rngs=key_dict)
```
I guess this would only work for custom-defined Modules.

> I'd personally suggest that flax should silently split up key inside the init and apply for me.

Passing in one key that will be used by all RNG streams by default, while still giving the user the option of explicitly passing in different RNG's for different RNG streams could be a good quality-of-life change. I'll discuss with the team more on this.
I'm not fully aware of the flax internals, I wasn't aware of `scope.rngs` so this is great to know! I can already use this to unit-test e.g., typos in creating rngs (e.g., I have the RNG stream `dropout` and somewhere else `drpout`).

Though in your code example you still have to pass in all keys explicitly in the init, which I might not know in advance for generic Modules and would have to explicitly trace.

It would be amazing if this would be handled internally. 

Resolving patch:
diff --git a/flax/core/scope.py b/flax/core/scope.py
--- a/flax/core/scope.py
+++ b/flax/core/scope.py
@@ -750,10 +750,13 @@ def has_rng(self, name: str) -> bool:
     """Returns true if a PRNGSequence with name `name` exists."""
     return name in self.rngs
 
-  def make_rng(self, name: str) -> PRNGKey:
+  def make_rng(self, name: str = 'default') -> PRNGKey:
     """Generates A PRNGKey from a PRNGSequence with name `name`."""
     if not self.has_rng(name):
-      raise errors.InvalidRngError(f'{self.name} needs PRNG for "{name}"')
+      if self.has_rng('default'):
+        name = 'default'
+      else:
+        raise errors.InvalidRngError(f'{self.name} needs PRNG for "{name}"')
     self._check_valid()
     self._validate_trace_level()
     self.rng_counters[name] += 1
diff --git a/flax/linen/module.py b/flax/linen/module.py
--- a/flax/linen/module.py
+++ b/flax/linen/module.py
@@ -1755,7 +1755,7 @@ def has_rng(self, name: str) -> bool:
       raise ValueError("Can't query for RNGs on unbound modules")
     return self.scope.has_rng(name)
 
-  def make_rng(self, name: str) -> KeyArray:
+  def make_rng(self, name: str = 'default') -> KeyArray:
     """Returns a new RNG key from a given RNG sequence for this Module.
 
     The new RNG key is split from the previous one. Thus, every call to


Test patch:
diff --git a/tests/linen/linen_module_test.py b/tests/linen/linen_module_test.py
--- a/tests/linen/linen_module_test.py
+++ b/tests/linen/linen_module_test.py
@@ -2522,6 +2522,145 @@ def __call__(self, x):
     self.assertIn('Bar_0/Dense_1', module_paths)
     self.assertEqual(type(module_paths['Bar_0/Dense_1']), nn.Dense)
 
+  def test_init_apply_default_rng(self):
+    class SubModel(nn.Module):
+      @nn.compact
+      def __call__(self, x, apply_dropout):
+        x = nn.Dense(8)(x)
+        x = nn.Dropout(0.8)(x, deterministic=not apply_dropout)
+        p = self.param(
+          'parameter', lambda key, shape: jax.random.normal(key, shape), x.shape
+        )
+        noise = jax.random.normal(self.make_rng('noise'), x.shape)
+        return x * p + noise
+
+    class Model(nn.Module):
+      @nn.compact
+      def __call__(self, x, apply_dropout):
+        x = nn.Dense(16)(x)
+        x = SubModel()(x, apply_dropout)
+        x = nn.Dropout(0.5)(x, deterministic=not apply_dropout)
+        v = self.variable(
+          'var_collection',
+          'variable',
+          lambda shape: jax.random.normal(self.make_rng('var_rng'), shape),
+          x.shape,
+        )
+        noise = jax.random.normal(self.make_rng('noise'), x.shape)
+        return x * v.value + noise
+
+    key0, key1, key2 = jax.random.split(jax.random.key(0), 3)
+    x = jax.random.normal(key0, (10, 4))
+    model = Model()
+
+    # test init equality
+    default_variables = model.init({'default': key1}, x, apply_dropout=False)
+    # adding 'default' rng shouldn't change anything
+    rngs = {'params': key1, 'var_rng': key1, 'noise': key1, 'default': key0}
+    explicit_variables = model.init(rngs, x, apply_dropout=False)
+    self.assertTrue(
+      jax.tree_util.tree_all(
+        jax.tree_map(
+          lambda v1, v2: (v1 == v2).all(), default_variables, explicit_variables
+        )
+      )
+    )
+
+    # test init inequality
+    rngs['default'] = key1  # adding 'default' rng shouldn't change anything
+    for rng_name in ('params', 'var_rng'):
+      rngs[rng_name] = key2
+      explicit_variables = model.init(rngs, x, apply_dropout=False)
+      self.assertFalse(
+        jax.tree_util.tree_all(
+          jax.tree_map(
+            lambda v1, v2: (v1 == v2).all(),
+            default_variables,
+            explicit_variables,
+          )
+        )
+      )
+      rngs[rng_name] = key1
+
+    # test apply equality
+    default_out = model.apply(
+      default_variables, x, apply_dropout=True, rngs={'default': key1}
+    )
+    # adding 'default' rng shouldn't change anything
+    rngs = {'dropout': key1, 'noise': key1, 'default': key0}
+    explicit_out = model.apply(
+      default_variables, x, apply_dropout=True, rngs=rngs
+    )
+    np.testing.assert_allclose(default_out, explicit_out)
+
+    # test apply inequality
+    rngs['default'] = key1  # adding 'default' rng shouldn't change anything
+    for rng_name in ('dropout', 'noise'):
+      rngs[rng_name] = key2
+      explicit_out = model.apply(
+        default_variables, x, apply_dropout=True, rngs=rngs
+      )
+      with self.assertRaises(AssertionError):
+        np.testing.assert_allclose(default_out, explicit_out, atol=1e-1)
+      rngs[rng_name] = key1
+
+  def test_default_make_rng(self):
+    class SubModel(nn.Module):
+      @nn.compact
+      def __call__(self, x):
+        noise = jax.random.normal(self.make_rng(), x.shape)
+        return x + noise
+
+    class Model(nn.Module):
+      @nn.compact
+      def __call__(self, x):
+        x = SubModel()(x)
+        noise = jax.random.normal(self.make_rng(), x.shape)
+        return x + noise
+
+    key0, key1 = jax.random.split(jax.random.key(0), 2)
+    x = jax.random.normal(key0, (10, 4))
+    default_out = Model().apply({}, x, rngs={'default': key1})
+
+    class SubModel(nn.Module):
+      @nn.compact
+      def __call__(self, x):
+        noise = jax.random.normal(self.make_rng('default'), x.shape)
+        return x + noise
+
+    class Model(nn.Module):
+      @nn.compact
+      def __call__(self, x):
+        x = SubModel()(x)
+        noise = jax.random.normal(self.make_rng('default'), x.shape)
+        return x + noise
+
+    explicit_out = Model().apply({}, x, rngs={'default': key1})
+    np.testing.assert_allclose(default_out, explicit_out)
+
+  def test_default_rng_error(self):
+    class Model(nn.Module):
+      @nn.compact
+      def __call__(self, x):
+        return nn.Dense(2)(x)
+
+    model = Model()
+    with self.assertRaisesRegex(
+      errors.InvalidRngError, 'Dense_0 needs PRNG for "params"'
+    ):
+      model.init({'other_rng_stream': jax.random.key(0)}, jnp.ones((1, 3)))
+
+    class Model(nn.Module):
+      @nn.compact
+      def __call__(self, x):
+        return x + jax.random.normal(self.make_rng('default'), x.shape)
+
+    model = Model()
+    with self.assertRaisesRegex(
+      errors.InvalidRngError, 'None needs PRNG for "default"'
+    ):
+      model.init(jax.random.key(0), jnp.ones((1, 3)))
+
 
 class LeakTests(absltest.TestCase):
   def test_tracer_leaks(self):