prompt: You will be provided with a partial code base and an issue statement explaining a problem to resolve.
  <issue>
  Make RNG name configurable in Dropout
  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:
  
  ```python
  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.
  
  </issue>
  <code>
  [start of flax/linen/stochastic.py]
  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     Note: When using :meth:`Module.apply() <flax.linen.Module.apply>`, make sure
  31     to include an RNG seed named `'dropout'`. For example::
  32 
  33       model.apply({'params': params}, inputs=inputs, train=True, rngs={'dropout': dropout_rng})`
  34 
  35     Attributes:
  36       rate: the dropout probability.  (_not_ the keep rate!)
  37       broadcast_dims: dimensions that will share the same dropout mask
  38       deterministic: if false the inputs are scaled by `1 / (1 - rate)` and
  39         masked, whereas if true, no mask is applied and the inputs are returned
  40         as is.
  41   """
  42   rate: float
  43   broadcast_dims: Sequence[int] = ()
  44   deterministic: Optional[bool] = None
  45 
  46   @compact
  47   def __call__(self, inputs, deterministic: Optional[bool] = None):
  48     """Applies a random dropout mask to the input.
  49 
  50     Args:
  51       inputs: the inputs that should be randomly masked.
  52       deterministic: if false the inputs are scaled by `1 / (1 - rate)` and
  53         masked, whereas if true, no mask is applied and the inputs are returned
  54         as is.
  55 
  56     Returns:
  57       The masked inputs reweighted to preserve mean.
  58     """
  59     deterministic = merge_param(
  60         'deterministic', self.deterministic, deterministic)
  61     if self.rate == 0.:
  62       return inputs
  63     # Prevent gradient NaNs in 1.0 edge-case.
  64     if self.rate == 1.0:
  65       return jnp.zeros_like(inputs)
  66     keep_prob = 1. - self.rate
  67     if deterministic:
  68       return inputs
  69     else:
  70       rng = self.make_rng('dropout')
  71       broadcast_shape = list(inputs.shape)
  72       for dim in self.broadcast_dims:
  73         broadcast_shape[dim] = 1
  74       mask = random.bernoulli(rng, p=keep_prob, shape=broadcast_shape)
  75       mask = jnp.broadcast_to(mask, inputs.shape)
  76       return lax.select(mask, inputs / keep_prob, jnp.zeros_like(inputs))
  77
  [end of flax/linen/stochastic.py]
  </code>
  
  I need you to solve this issue by generating a single patch file that I can apply directly to this repository using git apply. Please respond with a single patch file in the following format.
  <patch>
  diff --git a/file.py b/file.py
  --- a/file.py
  +++ b/file.py
  @@ -1,27 +1,35 @@
   def euclidean(a, b):
  -    while b:
  -        a, b = b, a % b
  -    return a
  +    if b == 0:
  +        return a
  +    return euclidean(b, a % b)
   
   
   def bresenham(x0, y0, x1, y1):
       points = []
       dx = abs(x1 - x0)
       dy = abs(y1 - y0)
  -    sx = 1 if x0 < x1 else -1
  -    sy = 1 if y0 < y1 else -1
  -    err = dx - dy
  +    x, y = x0, y0
  +    sx = -1 if x0 > x1 else 1
  +    sy = -1 if y0 > y1 else 1
   
  -    while True:
  -        points.append((x0, y0))
  -        if x0 == x1 and y0 == y1:
  -            break
  -        e2 = 2 * err
  -        if e2 > -dy:
  -            err -= dy
  -            x0 += sx
  -        if e2 < dx:
  -            err += dx
  -            y0 += sy
  +    if dx > dy:
  +        err = dx / 2.0
  +        while x != x1:
  +            points.append((x, y))
  +            err -= dy
  +            if err < 0:
  +                y += sy
  +                err += dx
  +            x += sx
  +    else:
  +        err = dy / 2.0
  +        while y != y1:
  +            points.append((x, y))
  +            err -= dx
  +            if err < 0:
  +                x += sx
  +                err += dy
  +            y += sy
  +
  +    points.append((x, y))
       return points
  
  </patch>
golden_diff: diff --git a/flax/linen/stochastic.py b/flax/linen/stochastic.py
  --- a/flax/linen/stochastic.py
  +++ b/flax/linen/stochastic.py
  @@ -38,10 +38,12 @@
         deterministic: if false the inputs are scaled by `1 / (1 - rate)` and
           masked, whereas if true, no mask is applied and the inputs are returned
           as is.
  +      rng_collection: the rng collection name to use when requesting an rng key.
     """
     rate: float
     broadcast_dims: Sequence[int] = ()
     deterministic: Optional[bool] = None
  +  rng_collection: str = 'dropout'
   
     @compact
     def __call__(self, inputs, deterministic: Optional[bool] = None):
  @@ -67,7 +69,7 @@
       if deterministic:
         return inputs
       else:
  -      rng = self.make_rng('dropout')
  +      rng = self.make_rng(self.rng_collection)
         broadcast_shape = list(inputs.shape)
         for dim in self.broadcast_dims:
           broadcast_shape[dim] = 1
  
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@@ -38,10 +38,12 @@\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+      rng_collection: the rng collection name to use when requesting an rng key.\n   \"\"\"\n   rate: float\n   broadcast_dims: Sequence[int] = ()\n   deterministic: Optional[bool] = None\n+  rng_collection: str = 'dropout'\n \n   @compact\n   def __call__(self, inputs, deterministic: Optional[bool] = None):\n@@ -67,7 +69,7 @@\n     if deterministic:\n       return inputs\n     else:\n-      rng = self.make_rng('dropout')\n+      rng = self.make_rng(self.rng_collection)\n       broadcast_shape = list(inputs.shape)\n       for dim in self.broadcast_dims:\n         broadcast_shape[dim] = 1\n", "issue": "Make RNG name configurable in Dropout\nCurrently `'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:\r\n1. A regular dropout which is applied to the inputs with a different mask at each step.\r\n2. A \"recurrent dropout\" that is applied to the state with the same mask at each step.\r\n\r\nTo 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:\r\n\r\n```python\r\nnn.scan(..., split_rngs={'dropout': True, 'recurrent_dropout': False})\r\n```\r\nThe 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.\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    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"}]}