Improve documentation for `Dropout` and `rngs` argument in `linen.Module.apply()`

[cccntu]

Here is an example of Dropout in a model definition: https://github.com/google/flax/blob/d068512a932da3e05b822790a591bac391aeab36/examples/nlp_seq/models.py#L211

Here is the apply(), where rngs is passed in https://github.com/google/flax/blob/d068512a932da3e05b822790a591bac391aeab36/examples/nlp_seq/train.py#L206-L207 However the rng is not very clearly explained in apply() https://github.com/google/flax/blob/615f40be774e7ed66fd344e8291ac0d48ebcef7d/flax/linen/module.py#L749 The rngs seems to be passed to flax/core/scope.py Here is the code for Dropout (linen) https://github.com/google/flax/blob/9b4807840c5cb26ef5e29028e3558d404aee00a0/flax/linen/stochastic.py#L56-L57 Here is the code for make_rng() https://github.com/google/flax/blob/615f40be774e7ed66fd344e8291ac0d48ebcef7d/flax/core/scope.py#L441-L447

The documentation for rngs in apply() should have a (pointer to) list of names of possible rngs And documentation for Dropout should mention how to pass in rng using apply(), without directly passing in like Dropout()(x,rng=rng). Also probably need to mention the make_rng() fold_in the rng so each dropout layer will use different rng if there are multiple dropout layers.

[andsteing]

We could mention that Dropout() requires an rng with the name dropout in its module documentation. The code is currently very short and it's easily visible, but I agree it would be better discoverable if it was mentioned in the class pydoc as well.

I also think that extending the Module.apply() could be extended with something like

    rngs: The rngs for the variable collections. For example :class:`flax.linen.stochastic.Dropout`
      requires an additional rng named `dropout`.

would be reasonable, since this is quite common and illustrative for newcomers.

WDYT @marcvanzee who was been thinking about the right level of verbosity in our docs

@cccntu why don't you give it a try updating the documentation and sending a PR?

[cccntu]

@andsteing Thanks for the suggestion. I will give it a try in a few days. :)