google / google/flax

Force no split in `make_rng`

Open
#3,115 5 comments 0 reactions 0 assignees View on GitHub
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

Open the contributing guide

Assessment

This issue has not been assessed yet.

Get new issues in your inbox

A short digest of beginner-friendly GitHub issues.