Make RNG name configurable in Dropout

[cgarciae]

Currently 'dropout' is hard-coded as an argument of make_rng inside Dropout. However, when implementing support for "recurrent dropout" in an LSTMCell or similar you need two kinds of dropout:

1. A regular dropout which is applied to the inputs with a different mask at each step.
2. A "recurrent dropout" that is applied to the state with the same mask at each step.

To implement 2 a possibility is to set the RNG name to 'recurrent_dropout' on the Dropout layer applied to the state and guarantee that each step uses the same random state. From nn.scan's perspective the would look like this:

nn.scan(..., split_rngs={'dropout': True, 'recurrent_dropout': False})

The proposal is to add an rng_name (or similar) attribute to Dropout so we are able support these kind of use-cases. The alternative would be to create a separate RecurrentDropout layer with the same code but different hard-coded value.

[jheek]

It might be more efficient to compute the mask outside of the loop if you are going to share it. random.* operations are usually quite expensive compared to other operations

[cgarciae]

Its a good point. Maybe we need a RNNCellBase.get_stochastic_mask API and have cells optionally accept a stochastic_mask argument.

[cgarciae]

Some updates: creating a mask API for RNNs resulted in more complex code. Leveraging RNG collections with transformation like nn.scan results in very understandable code and effectively uses Flax existing APIs, if performance becomes an issue then we should just advice the users to use the jax_default_prng_impl flag for faster alternatives as stated in Advanced RNG configuration.