Force no split in `make_rng`
- Dominant language
- Jupyter Notebook
- Stars
- 7.3k
- Forks
- 833
- Avg merge
- 5h 11m
- Merged PRs (30d)
- 5
Description
### Discussed in https://github.com/google/flax/discussions/3113
Originally posted by **zaccharieramzi** May 24, 2023
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.
Contributor guide
Assessment
This issue has not been assessed yet.