Repository: google/flax

Issue:
Force no split in `make_rng`
### Discussed in https://github.com/google/flax/discussions/3113

<div type='discussions-op-text'>

<sup>Originally posted by **zaccharieramzi** May 24, 2023</sup>
I have the following situation: I am using a `Dropout` layer multiple times without a [`nn.scan` or `nn.while_loop`](https://github.com/google/flax/discussions/2920#discussioncomment-5180446), therefore I cannot use `split_rngs={"dropout": False}`. 
However, I would still like to use the same dropout mask twice.

Is it possible to specify "no split" to make rng for certain collections?

If I just take the original dropout example I would like to do something like:

```python
# Setup.
import jax
import jax.numpy as jnp
import flax.linen as nn

# Randomness.
seed = 0
root_key = jax.random.PRNGKey(seed=seed)
main_key, params_key, dropout_key = jax.random.split(key=root_key, num=3)

# A simple network.
class MyModel(nn.Module):
  num_neurons: int
  training: bool
  @nn.compact
  def __call__(self, x):
    x = nn.Dense(self.num_neurons)(x)
    # Set the dropout layer with a rate of 50% .
    # When the `deterministic` flag is `True`, dropout is turned off.
    x = nn.Dropout(rate=0.5, deterministic=not self.training)(x)
    x = nn.Dropout(rate=0.5, deterministic=not self.training)(x)
    return x

# Instantiate `MyModel` (you don't need to set `training=True` to
# avoid performing the forward pass computation).
my_model = MyModel(num_neurons=3, training=False)

x = jax.random.uniform(key=main_key, shape=(3, 4, 4))

# Initialize with `flax.linen.init()`.
# The `params_key` is equivalent to a dictionary of PRNGs.
# (Here, you are providing only one PRNG key.) 
variables = my_model.init(params_key, x)

# Perform the forward pass with `flax.linen.apply()`.
my_model.training = True
y = my_model.apply(variables, x, rngs={'dropout': dropout_key})
```

and still have `jnp.sum(y == 0.) / (3*4*3) == 0.5` approx.


For more context I am actually trying to implement [Deep Equilibrium Models](https://arxiv.org/abs/1909.01377) using [`jaxopt`](https://jaxopt.github.io/stable/) and `flax`, where the fixed point defining function uses dropout. 
I also [tried to see](https://github.com/google/jaxopt/issues/432) if the `split_rngs` functionality could be extended to `jaxopt` but I think it's going to be difficult.</div>


Discussion:
Hey @zaccharieramzi, I've converted the discussion into and issue as it seems something that we should improve.
I've created #3114, which would allow you to optionally specify the `rng` key for each `Dropout` layer, e.g:

```python
# A simple network.
class MyModel(nn.Module):
  num_neurons: int
  training: bool
  @nn.compact
  def __call__(self, x):
    x = nn.Dense(self.num_neurons)(x)
    # Set the dropout layer with a rate of 50% .
    # When the `deterministic` flag is `True`, dropout is turned off.
    key = self.make_rng('dropout')
    x = nn.Dropout(rate=0.5, deterministic=not self.training)(x, rng=key)
    x = nn.Dropout(rate=0.5, deterministic=not self.training)(x, rng=key)
    return 
```
This way both layers will produce the same mask.

Resolving patch:
diff --git a/flax/configurations.py b/flax/configurations.py
--- a/flax/configurations.py
+++ b/flax/configurations.py
@@ -12,7 +12,7 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 
-"""Global configuration options for Flax.
+r"""Global configuration options for Flax.
 
 Now a wrapper over jax.config, in which all config vars have a 'flax\_' prefix.
 
diff --git a/flax/linen/module.py b/flax/linen/module.py
--- a/flax/linen/module.py
+++ b/flax/linen/module.py
@@ -49,8 +49,8 @@
 
 traceback_util.register_exclusion(__file__)
 
-PRNGKey = Any  # pylint: disable=invalid-name
-RNGSequences = Dict[str, PRNGKey]
+KeyArray = Union[jax.Array, jax.random.KeyArray]  # pylint: disable=invalid-name
+RNGSequences = Dict[str, KeyArray]
 Array = Any    # pylint: disable=invalid-name
 
 
@@ -1270,7 +1270,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) -> PRNGKey:
+  def make_rng(self, name: str) -> 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
@@ -1494,7 +1494,7 @@ def other_fn(instance, ...):
 
   @traceback_util.api_boundary
   def init_with_output(self,
-                       rngs: Union[PRNGKey, RNGSequences],
+                       rngs: Union[KeyArray, RNGSequences],
                        *args,
                        method: Union[Callable[..., Any], str, None] = None,
                        mutable: CollectionFilter = DenyList('intermediates'),
@@ -1551,7 +1551,7 @@ def init_with_output(self,
 
   @traceback_util.api_boundary
   def init(self,
-           rngs: Union[PRNGKey, RNGSequences],
+           rngs: Union[KeyArray, RNGSequences],
            *args,
            method: Union[Callable[..., Any], str, None] = None,
            mutable: CollectionFilter = DenyList('intermediates'),
@@ -1648,7 +1648,7 @@ def init(self,
 
   @traceback_util.api_boundary
   def lazy_init(self,
-           rngs: Union[PRNGKey, RNGSequences],
+           rngs: Union[KeyArray, RNGSequences],
            *args,
            method: Optional[Callable[..., Any]] = None,
            mutable: CollectionFilter = DenyList('intermediates'),
@@ -1875,7 +1875,7 @@ def _root_has_collection():
 
   def tabulate(
     self,
-    rngs: Union[PRNGKey, RNGSequences],
+    rngs: Union[KeyArray, RNGSequences],
     *args,
     depth: Optional[int] = None,
     show_repeated: bool = False,
diff --git a/flax/linen/stochastic.py b/flax/linen/stochastic.py
--- a/flax/linen/stochastic.py
+++ b/flax/linen/stochastic.py
@@ -14,7 +14,9 @@
 
 """Stochastic modules."""
 
-from typing import Optional, Sequence
+from typing import Optional, Sequence, Union
+
+import jax
 
 from flax.linen.module import compact
 from flax.linen.module import merge_param
@@ -23,6 +25,8 @@
 from jax import random
 import jax.numpy as jnp
 
+KeyArray = Union[jax.Array, jax.random.KeyArray]
+
 
 class Dropout(Module):
   """Create a dropout layer.
@@ -46,7 +50,9 @@ class Dropout(Module):
   rng_collection: str = 'dropout'
 
   @compact
-  def __call__(self, inputs, deterministic: Optional[bool] = None):
+  def __call__(
+    self, inputs, deterministic: Optional[bool] = None, rng: Optional[KeyArray] = None
+  ):
     """Applies a random dropout mask to the input.
 
     Args:
@@ -54,6 +60,8 @@ def __call__(self, inputs, deterministic: Optional[bool] = None):
       deterministic: if false the inputs are scaled by `1 / (1 - rate)` and
         masked, whereas if true, no mask is applied and the inputs are returned
         as is.
+      rng: an optional PRNGKey used as the random key, if not specified, one
+        will be generated using ``make_rng`` with the ``rng_collection`` name.
 
     Returns:
       The masked inputs reweighted to preserve mean.
@@ -69,7 +77,8 @@ def __call__(self, inputs, deterministic: Optional[bool] = None):
       return jnp.zeros_like(inputs)
 
     keep_prob = 1. - self.rate
-    rng = self.make_rng(self.rng_collection)
+    if rng is None:
+      rng = self.make_rng(self.rng_collection)
     broadcast_shape = list(inputs.shape)
     for dim in self.broadcast_dims:
       broadcast_shape[dim] = 1


Test patch:
diff --git a/tests/linen/linen_test.py b/tests/linen/linen_test.py
--- a/tests/linen/linen_test.py
+++ b/tests/linen/linen_test.py
@@ -323,6 +323,21 @@ def test_dropout_rate_limits(self):
     res = jax.grad(lambda x, k: jnp.sum(fn(x, k)))(inputs, key3)
     self.assertFalse(np.isnan(res).any())
 
+  def test_dropout_manual_rng(self):
+    class Foo(nn.Module):
+      @nn.compact
+      def __call__(self, x):
+        key = self.make_rng('dropout')
+        x1 = nn.Dropout(rate=0.5, deterministic=False)(x, rng=key)
+        x2 = nn.Dropout(rate=0.5, deterministic=False)(x, rng=key)
+        return x1, x2
+
+    module = Foo()
+    x1, x2 = module.apply(
+      {}, jnp.ones((20, 20)), rngs={'dropout': random.PRNGKey(0)})
+
+    np.testing.assert_array_equal(x1, x2)
+
 
 # TODO(flax-dev): add integration tests for RNN cells
 class RecurrentTest(absltest.TestCase):