prompt: You will be provided with a partial code base and an issue statement explaining a problem to resolve.
  <issue>
  PReLU activation implementation
  I wanted to gauge interest on adding a PReLU activation. I noticed that `flax.linen.activations` are simply aliasing `jax.nn` activation functions which also doesn't have a PReLU implementation.
  
  To add some background, PReLU is simply Leaky ReLU where the alpha (slope) parameter is trainable and not fixed. This makes it simple to implement as a Module if desired.
  
  Here's an example implementation from another [project](https://github.com/isaaccorley/jax-enhance) of mine. 
  
  ```python
  from functools import partial
  from typing import Any, Sequence
  
  import jax.numpy as jnp
  import flax.linen as nn
  
  
  # This is nearly identical to jnp.ones however multiplies the output of jnp.ones by the constant value
  def constant(key, shape: Sequence[int], value: Any, dtype: Any = jnp.float32) -> jnp.ndarray:
      value = jnp.asarray(value, dtype)
      return jnp.ones(shape, dtype) * value
  
  
  class PReLU(nn.Module):
      negative_slope_init: float = 0.01
      dtype: Any = jnp.float32
  
      @nn.compact
      def __call__(self, x: jnp.ndarray) -> jnp.ndarray:
          x = jnp.asarray(x, self.dtype)
          negative_slope = self.param(
              "negative_slope",
              partial(constant, value=self.negative_slope_init, dtype=self.dtype),
              (1,)
          )
          return jnp.where(x >= 0, x, negative_slope * x)
  ```
  
  </issue>
  <code>
  [start of flax/linen/__init__.py]
  1 # Copyright 2021 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 """The Flax Module system."""
  16 
  17 
  18 # pylint: disable=g-multiple-import
  19 # re-export commonly used modules and functions
  20 from .activation import (celu, elu, gelu, glu, leaky_relu, log_sigmoid,
  21                          log_softmax, relu, sigmoid, soft_sign, softmax,
  22                          softplus, swish, silu, tanh)
  23 from .attention import (MultiHeadDotProductAttention, SelfAttention,
  24                         dot_product_attention, make_attention_mask,
  25                         make_causal_mask, combine_masks)
  26 from ..core import broadcast, DenyList, FrozenDict
  27 from .linear import Conv, ConvTranspose, Dense, DenseGeneral, Embed
  28 from .module import (Module, compact, nowrap, enable_named_call,
  29                      disable_named_call, override_named_call, Variable, init,
  30                      init_with_output, apply, merge_param)
  31 from .normalization import BatchNorm, GroupNorm, LayerNorm
  32 from .pooling import avg_pool, max_pool
  33 from .recurrent import GRUCell, LSTMCell, ConvLSTM, OptimizedLSTMCell
  34 from .stochastic import Dropout
  35 from .transforms import jit, named_call, checkpoint, remat, remat_scan, scan, vmap, map_variables, vjp, jvp
  36 from .initializers import zeros, ones
  37 
  38 # pylint: enable=g-multiple-import
  39
  [end of flax/linen/__init__.py]
  [start of flax/linen/activation.py]
  1 # Copyright 2021 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 """Activation functions.
  16 """
  17 
  18 # pylint: disable=unused-import
  19 # re-export activation functions from jax.nn
  20 from jax.nn import celu
  21 from jax.nn import elu
  22 from jax.nn import gelu
  23 from jax.nn import glu
  24 from jax.nn import leaky_relu
  25 from jax.nn import log_sigmoid
  26 from jax.nn import log_softmax
  27 from jax.nn import normalize
  28 from jax.nn import relu
  29 from jax.nn import sigmoid
  30 from jax.nn import soft_sign
  31 from jax.nn import softmax
  32 from jax.nn import softplus
  33 from jax.nn import swish
  34 from jax.nn import silu
  35 from jax.nn import selu
  36 from jax.nn import hard_tanh
  37 from jax.nn import relu6
  38 from jax.nn import hard_sigmoid
  39 from jax.nn import hard_swish
  40 
  41 from jax.numpy import tanh
  42 # pylint: enable=unused-import
  43
  [end of flax/linen/activation.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/__init__.py b/flax/linen/__init__.py
  --- a/flax/linen/__init__.py
  +++ b/flax/linen/__init__.py
  @@ -19,7 +19,7 @@
   # re-export commonly used modules and functions
   from .activation import (celu, elu, gelu, glu, leaky_relu, log_sigmoid,
                            log_softmax, relu, sigmoid, soft_sign, softmax,
  -                         softplus, swish, silu, tanh)
  +                         softplus, swish, silu, tanh, PReLU)
   from .attention import (MultiHeadDotProductAttention, SelfAttention,
                           dot_product_attention, make_attention_mask,
                           make_causal_mask, combine_masks)
  diff --git a/flax/linen/activation.py b/flax/linen/activation.py
  --- a/flax/linen/activation.py
  +++ b/flax/linen/activation.py
  @@ -40,3 +40,35 @@
   
   from jax.numpy import tanh
   # pylint: enable=unused-import
  +
  +from typing import Any
  +
  +from flax.linen.module import Module, compact
  +import jax.numpy as jnp
  +
  +
  +Array = Any
  +
  +
  +class PReLU(Module):
  +  """Parametric Rectified Linear Unit (PReLU) activation function.
  +
  +  Attributes:
  +    negative_slope_init: the value to initialize the negative slope.
  +  """
  +  negative_slope_init: float = 0.01
  +  @compact
  +  def __call__(self, inputs: Array) -> Array:
  +    """Applies an activation to the inputs.
  +
  +    Args:
  +      inputs: the nd-array to apply the activation function to.
  +
  +    Returns:
  +      The transformed input.
  +    """
  +    negative_slope = self.param(
  +      'negative_slope',
  +      lambda k: jnp.asarray(self.negative_slope_init, jnp.float32)
  +    )
  +    return jnp.where(inputs >= 0, inputs, jnp.asarray(negative_slope, inputs.dtype) * inputs)
  
verification_info: {"golden_diff": "diff --git a/flax/linen/__init__.py b/flax/linen/__init__.py\n--- a/flax/linen/__init__.py\n+++ b/flax/linen/__init__.py\n@@ -19,7 +19,7 @@\n # re-export commonly used modules and functions\n from .activation import (celu, elu, gelu, glu, leaky_relu, log_sigmoid,\n                          log_softmax, relu, sigmoid, soft_sign, softmax,\n-                         softplus, swish, silu, tanh)\n+                         softplus, swish, silu, tanh, PReLU)\n from .attention import (MultiHeadDotProductAttention, SelfAttention,\n                         dot_product_attention, make_attention_mask,\n                         make_causal_mask, combine_masks)\ndiff --git a/flax/linen/activation.py b/flax/linen/activation.py\n--- a/flax/linen/activation.py\n+++ b/flax/linen/activation.py\n@@ -40,3 +40,35 @@\n \n from jax.numpy import tanh\n # pylint: enable=unused-import\n+\n+from typing import Any\n+\n+from flax.linen.module import Module, compact\n+import jax.numpy as jnp\n+\n+\n+Array = Any\n+\n+\n+class PReLU(Module):\n+  \"\"\"Parametric Rectified Linear Unit (PReLU) activation function.\n+\n+  Attributes:\n+    negative_slope_init: the value to initialize the negative slope.\n+  \"\"\"\n+  negative_slope_init: float = 0.01\n+  @compact\n+  def __call__(self, inputs: Array) -> Array:\n+    \"\"\"Applies an activation to the inputs.\n+\n+    Args:\n+      inputs: the nd-array to apply the activation function to.\n+\n+    Returns:\n+      The transformed input.\n+    \"\"\"\n+    negative_slope = self.param(\n+      'negative_slope',\n+      lambda k: jnp.asarray(self.negative_slope_init, jnp.float32)\n+    )\n+    return jnp.where(inputs >= 0, inputs, jnp.asarray(negative_slope, inputs.dtype) * inputs)\n", "issue": "PReLU activation implementation\nI wanted to gauge interest on adding a PReLU activation. I noticed that `flax.linen.activations` are simply aliasing `jax.nn` activation functions which also doesn't have a PReLU implementation.\r\n\r\nTo add some background, PReLU is simply Leaky ReLU where the alpha (slope) parameter is trainable and not fixed. This makes it simple to implement as a Module if desired.\r\n\r\nHere's an example implementation from another [project](https://github.com/isaaccorley/jax-enhance) of mine. \r\n\r\n```python\r\nfrom functools import partial\r\nfrom typing import Any, Sequence\r\n\r\nimport jax.numpy as jnp\r\nimport flax.linen as nn\r\n\r\n\r\n# This is nearly identical to jnp.ones however multiplies the output of jnp.ones by the constant value\r\ndef constant(key, shape: Sequence[int], value: Any, dtype: Any = jnp.float32) -> jnp.ndarray:\r\n    value = jnp.asarray(value, dtype)\r\n    return jnp.ones(shape, dtype) * value\r\n\r\n\r\nclass PReLU(nn.Module):\r\n    negative_slope_init: float = 0.01\r\n    dtype: Any = jnp.float32\r\n\r\n    @nn.compact\r\n    def __call__(self, x: jnp.ndarray) -> jnp.ndarray:\r\n        x = jnp.asarray(x, self.dtype)\r\n        negative_slope = self.param(\r\n            \"negative_slope\",\r\n            partial(constant, value=self.negative_slope_init, dtype=self.dtype),\r\n            (1,)\r\n        )\r\n        return jnp.where(x >= 0, x, negative_slope * x)\r\n```\n", "before_files": [{"content": "# Copyright 2021 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\"\"\"The Flax Module system.\"\"\"\n\n\n# pylint: disable=g-multiple-import\n# re-export commonly used modules and functions\nfrom .activation import (celu, elu, gelu, glu, leaky_relu, log_sigmoid,\n                         log_softmax, relu, sigmoid, soft_sign, softmax,\n                         softplus, swish, silu, tanh)\nfrom .attention import (MultiHeadDotProductAttention, SelfAttention,\n                        dot_product_attention, make_attention_mask,\n                        make_causal_mask, combine_masks)\nfrom ..core import broadcast, DenyList, FrozenDict\nfrom .linear import Conv, ConvTranspose, Dense, DenseGeneral, Embed\nfrom .module import (Module, compact, nowrap, enable_named_call,\n                     disable_named_call, override_named_call, Variable, init,\n                     init_with_output, apply, merge_param)\nfrom .normalization import BatchNorm, GroupNorm, LayerNorm\nfrom .pooling import avg_pool, max_pool\nfrom .recurrent import GRUCell, LSTMCell, ConvLSTM, OptimizedLSTMCell\nfrom .stochastic import Dropout\nfrom .transforms import jit, named_call, checkpoint, remat, remat_scan, scan, vmap, map_variables, vjp, jvp\nfrom .initializers import zeros, ones\n\n# pylint: enable=g-multiple-import\n", "path": "flax/linen/__init__.py"}, {"content": "# Copyright 2021 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\"\"\"Activation functions.\n\"\"\"\n\n# pylint: disable=unused-import\n# re-export activation functions from jax.nn\nfrom jax.nn import celu\nfrom jax.nn import elu\nfrom jax.nn import gelu\nfrom jax.nn import glu\nfrom jax.nn import leaky_relu\nfrom jax.nn import log_sigmoid\nfrom jax.nn import log_softmax\nfrom jax.nn import normalize\nfrom jax.nn import relu\nfrom jax.nn import sigmoid\nfrom jax.nn import soft_sign\nfrom jax.nn import softmax\nfrom jax.nn import softplus\nfrom jax.nn import swish\nfrom jax.nn import silu\nfrom jax.nn import selu\nfrom jax.nn import hard_tanh\nfrom jax.nn import relu6\nfrom jax.nn import hard_sigmoid\nfrom jax.nn import hard_swish\n\nfrom jax.numpy import tanh\n# pylint: enable=unused-import\n", "path": "flax/linen/activation.py"}]}