prompt: We are currently solving the following issue within our repository. Here is the issue text:
  --- BEGIN ISSUE ---
  Improve documentation for `Dropout` and `rngs` argument in `linen.Module.apply()`
  
  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.
  
  --- END ISSUE ---
  
  Below are some code segments, each from a relevant file. One or more of these files may contain bugs.
  --- BEGIN FILES ---
  Path: `flax/linen/stochastic.py`
  Content:
  ```
  1 # Copyright 2022 The Flax Authors.
  2 #
  3 # Licensed under the Apache License, Version 2.0 (the "License");
  4 # you may not use this file except in compliance with the License.
  5 # You may obtain a copy of the License at
  6 #
  7 #     http://www.apache.org/licenses/LICENSE-2.0
  8 #
  9 # Unless required by applicable law or agreed to in writing, software
  10 # distributed under the License is distributed on an "AS IS" BASIS,
  11 # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
  12 # See the License for the specific language governing permissions and
  13 # limitations under the License.
  14 
  15 """Stochastic modules."""
  16 
  17 from typing import Optional, Sequence
  18 
  19 from flax.linen.module import compact
  20 from flax.linen.module import merge_param
  21 from flax.linen.module import Module
  22 from jax import lax
  23 from jax import random
  24 import jax.numpy as jnp
  25 
  26 
  27 class Dropout(Module):
  28   """Create a dropout layer.
  29 
  30     Attributes:
  31       rate: the dropout probability.  (_not_ the keep rate!)
  32       broadcast_dims: dimensions that will share the same dropout mask
  33       deterministic: if false the inputs are scaled by `1 / (1 - rate)` and
  34         masked, whereas if true, no mask is applied and the inputs are returned
  35         as is.
  36   """
  37   rate: float
  38   broadcast_dims: Sequence[int] = ()
  39   deterministic: Optional[bool] = None
  40 
  41   @compact
  42   def __call__(self, inputs, deterministic: Optional[bool] = None):
  43     """Applies a random dropout mask to the input.
  44 
  45     Args:
  46       inputs: the inputs that should be randomly masked.
  47       deterministic: if false the inputs are scaled by `1 / (1 - rate)` and
  48         masked, whereas if true, no mask is applied and the inputs are returned
  49         as is.
  50 
  51     Returns:
  52       The masked inputs reweighted to preserve mean.
  53     """
  54     deterministic = merge_param(
  55         'deterministic', self.deterministic, deterministic)
  56     if self.rate == 0.:
  57       return inputs
  58     # Prevent gradient NaNs in 1.0 edge-case.
  59     if self.rate == 1.0:
  60       return jnp.zeros_like(inputs)
  61     keep_prob = 1. - self.rate
  62     if deterministic:
  63       return inputs
  64     else:
  65       rng = self.make_rng('dropout')
  66       broadcast_shape = list(inputs.shape)
  67       for dim in self.broadcast_dims:
  68         broadcast_shape[dim] = 1
  69       mask = random.bernoulli(rng, p=keep_prob, shape=broadcast_shape)
  70       mask = jnp.broadcast_to(mask, inputs.shape)
  71       return lax.select(mask, inputs / keep_prob, jnp.zeros_like(inputs))
  72
  ```
  --- END FILES ---
  
  Please first localize the bug based on the issue statement, and then generate a patch according to the `git diff` format fenced by three backticks.
  
  Here is an example:
  
  ```diff
  diff --git a/examples/server_async.py b/examples/server_async.py
  --- a/examples/server_async.py
  +++ b/examples/server_async.py
  @@ -313,4 +313,4 @@
   
   
   if __name__ == "__main__":
  -    asyncio.run(run_async_server("."), debug=True)
  +    asyncio.run(run_async_server(), debug=True)
  diff --git a/examples/server_sync.py b/examples/server_sync.py
  --- a/examples/server_sync.py
  +++ b/examples/server_sync.py
  @@ -313,5 +313,5 @@
   
   
   if __name__ == "__main__":
  -    server = run_sync_server(".")
  +    server = run_sync_server()
       server.shutdown()
  
  ```
  
golden_diff: diff --git a/flax/linen/stochastic.py b/flax/linen/stochastic.py
  --- a/flax/linen/stochastic.py
  +++ b/flax/linen/stochastic.py
  @@ -27,6 +27,11 @@
   class Dropout(Module):
     """Create a dropout layer.
   
  +    Note: When using :meth:`Module.apply() <flax.linen.Module.apply>`, make sure
  +    to include an RNG seed named `'dropout'`. For example::
  +    
  +      model.apply({'params': params}, inputs=inputs, train=True, rngs={'dropout': dropout_rng})`
  +
       Attributes:
         rate: the dropout probability.  (_not_ the keep rate!)
         broadcast_dims: dimensions that will share the same dropout mask
  
verification_info: {"golden_diff": "diff --git a/flax/linen/stochastic.py b/flax/linen/stochastic.py\n--- a/flax/linen/stochastic.py\n+++ b/flax/linen/stochastic.py\n@@ -27,6 +27,11 @@\n class Dropout(Module):\n   \"\"\"Create a dropout layer.\n \n+    Note: When using :meth:`Module.apply() <flax.linen.Module.apply>`, make sure\n+    to include an RNG seed named `'dropout'`. For example::\n+    \n+      model.apply({'params': params}, inputs=inputs, train=True, rngs={'dropout': dropout_rng})`\n+\n     Attributes:\n       rate: the dropout probability.  (_not_ the keep rate!)\n       broadcast_dims: dimensions that will share the same dropout mask\n", "issue": "Improve documentation for `Dropout` and `rngs` argument in `linen.Module.apply()`\n\r\nHere is an example of `Dropout` in a model definition:\r\nhttps://github.com/google/flax/blob/d068512a932da3e05b822790a591bac391aeab36/examples/nlp_seq/models.py#L211\r\n\r\nHere is the `apply()`, where `rngs` is passed in\r\nhttps://github.com/google/flax/blob/d068512a932da3e05b822790a591bac391aeab36/examples/nlp_seq/train.py#L206-L207\r\nHowever the `rng` is not very clearly explained in `apply()`\r\nhttps://github.com/google/flax/blob/615f40be774e7ed66fd344e8291ac0d48ebcef7d/flax/linen/module.py#L749\r\nThe `rngs` seems to be passed to `flax/core/scope.py`\r\nHere is the code for `Dropout` (linen)\r\nhttps://github.com/google/flax/blob/9b4807840c5cb26ef5e29028e3558d404aee00a0/flax/linen/stochastic.py#L56-L57\r\nHere is the code for  `make_rng()`\r\nhttps://github.com/google/flax/blob/615f40be774e7ed66fd344e8291ac0d48ebcef7d/flax/core/scope.py#L441-L447\r\n\r\nThe documentation for `rngs` in `apply()` should have a (pointer to) list of names of possible rngs\r\nAnd documentation for `Dropout` should mention how to pass in rng using `apply()`, without directly passing in like `Dropout()(x,rng=rng)`.\r\nAlso 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.\n", "before_files": [{"content": "# Copyright 2022 The Flax Authors.\n#\n# Licensed under the Apache License, Version 2.0 (the \"License\");\n# you may not use this file except in compliance with the License.\n# You may obtain a copy of the License at\n#\n#     http://www.apache.org/licenses/LICENSE-2.0\n#\n# Unless required by applicable law or agreed to in writing, software\n# distributed under the License is distributed on an \"AS IS\" BASIS,\n# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.\n# See the License for the specific language governing permissions and\n# limitations under the License.\n\n\"\"\"Stochastic modules.\"\"\"\n\nfrom typing import Optional, Sequence\n\nfrom flax.linen.module import compact\nfrom flax.linen.module import merge_param\nfrom flax.linen.module import Module\nfrom jax import lax\nfrom jax import random\nimport jax.numpy as jnp\n\n\nclass Dropout(Module):\n  \"\"\"Create a dropout layer.\n\n    Attributes:\n      rate: the dropout probability.  (_not_ the keep rate!)\n      broadcast_dims: dimensions that will share the same dropout mask\n      deterministic: if false the inputs are scaled by `1 / (1 - rate)` and\n        masked, whereas if true, no mask is applied and the inputs are returned\n        as is.\n  \"\"\"\n  rate: float\n  broadcast_dims: Sequence[int] = ()\n  deterministic: Optional[bool] = None\n\n  @compact\n  def __call__(self, inputs, deterministic: Optional[bool] = None):\n    \"\"\"Applies a random dropout mask to the input.\n\n    Args:\n      inputs: the inputs that should be randomly masked.\n      deterministic: if false the inputs are scaled by `1 / (1 - rate)` and\n        masked, whereas if true, no mask is applied and the inputs are returned\n        as is.\n\n    Returns:\n      The masked inputs reweighted to preserve mean.\n    \"\"\"\n    deterministic = merge_param(\n        'deterministic', self.deterministic, deterministic)\n    if self.rate == 0.:\n      return inputs\n    # Prevent gradient NaNs in 1.0 edge-case.\n    if self.rate == 1.0:\n      return jnp.zeros_like(inputs)\n    keep_prob = 1. - self.rate\n    if deterministic:\n      return inputs\n    else:\n      rng = self.make_rng('dropout')\n      broadcast_shape = list(inputs.shape)\n      for dim in self.broadcast_dims:\n        broadcast_shape[dim] = 1\n      mask = random.bernoulli(rng, p=keep_prob, shape=broadcast_shape)\n      mask = jnp.broadcast_to(mask, inputs.shape)\n      return lax.select(mask, inputs / keep_prob, jnp.zeros_like(inputs))\n", "path": "flax/linen/stochastic.py"}], "after_files": [{"content": "# Copyright 2022 The Flax Authors.\n#\n# Licensed under the Apache License, Version 2.0 (the \"License\");\n# you may not use this file except in compliance with the License.\n# You may obtain a copy of the License at\n#\n#     http://www.apache.org/licenses/LICENSE-2.0\n#\n# Unless required by applicable law or agreed to in writing, software\n# distributed under the License is distributed on an \"AS IS\" BASIS,\n# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.\n# See the License for the specific language governing permissions and\n# limitations under the License.\n\n\"\"\"Stochastic modules.\"\"\"\n\nfrom typing import Optional, Sequence\n\nfrom flax.linen.module import compact\nfrom flax.linen.module import merge_param\nfrom flax.linen.module import Module\nfrom jax import lax\nfrom jax import random\nimport jax.numpy as jnp\n\n\nclass Dropout(Module):\n  \"\"\"Create a dropout layer.\n\n    Note: When using :meth:`Module.apply() <flax.linen.Module.apply>`, make sure\n    to include an RNG seed named `'dropout'`. For example::\n    \n      model.apply({'params': params}, inputs=inputs, train=True, rngs={'dropout': dropout_rng})`\n\n    Attributes:\n      rate: the dropout probability.  (_not_ the keep rate!)\n      broadcast_dims: dimensions that will share the same dropout mask\n      deterministic: if false the inputs are scaled by `1 / (1 - rate)` and\n        masked, whereas if true, no mask is applied and the inputs are returned\n        as is.\n  \"\"\"\n  rate: float\n  broadcast_dims: Sequence[int] = ()\n  deterministic: Optional[bool] = None\n\n  @compact\n  def __call__(self, inputs, deterministic: Optional[bool] = None):\n    \"\"\"Applies a random dropout mask to the input.\n\n    Args:\n      inputs: the inputs that should be randomly masked.\n      deterministic: if false the inputs are scaled by `1 / (1 - rate)` and\n        masked, whereas if true, no mask is applied and the inputs are returned\n        as is.\n\n    Returns:\n      The masked inputs reweighted to preserve mean.\n    \"\"\"\n    deterministic = merge_param(\n        'deterministic', self.deterministic, deterministic)\n    if self.rate == 0.:\n      return inputs\n    # Prevent gradient NaNs in 1.0 edge-case.\n    if self.rate == 1.0:\n      return jnp.zeros_like(inputs)\n    keep_prob = 1. - self.rate\n    if deterministic:\n      return inputs\n    else:\n      rng = self.make_rng('dropout')\n      broadcast_shape = list(inputs.shape)\n      for dim in self.broadcast_dims:\n        broadcast_shape[dim] = 1\n      mask = random.bernoulli(rng, p=keep_prob, shape=broadcast_shape)\n      mask = jnp.broadcast_to(mask, inputs.shape)\n      return lax.select(mask, inputs / keep_prob, jnp.zeros_like(inputs))\n", "path": "flax/linen/stochastic.py"}]}