We are currently solving the following issue within our repository. Here is the issue text:
--- BEGIN 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)
```

--- 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/__init__.py`
Content:
```
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
```

Path: `flax/linen/activation.py`
Content:
```
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 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()

```

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)

{"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"}], "after_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, PReLU)\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\nfrom typing import Any\n\nfrom flax.linen.module import Module, compact\nimport jax.numpy as jnp\n\n\nArray = Any\n\n\nclass 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", "path": "flax/linen/activation.py"}]}