Repository: google/flax

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)
```


Discussion:
Given that all current activation functions reside in JAX, it seem more fitting to add this JAX. Do you want to file an issue against their repo?
Thanks for the suggestion. The main reason I filed the issue here was because it seems like PReLU is a special case where it has a trainable param and, if I'm not mistaken, all other jax activations do not.

I'm not sure if this changes your suggestion, but it's something to consider.
@isaaccorley - hey so sorry for the slow feedback on your suggestion here.

2 points:
 - instead of defining a constant init func, we can just declare a jnp scalar array of the correct dtype. 
 - I think an -activation- "function" should strictly follow the dtype of its argument, so no dtype attribute, just derive it from `x`

So what if we added something like this?

```python
class PReLU(nn.Module):
    negative_slope_init: float = 0.01
    @nn.compact
    def __call__(self, x: jnp.ndarray) -> jnp.ndarray:
        negative_slope = self.param(
            "negative_slope",
            lambda k: jnp.array(self.negative_slope_init, x.dtype)
        )
        return jnp.where(x >= 0, x, negative_slope * x)
```
I'm indifferent on the implementation. I think the only thing to point out would be since we are inheriting from Module and other Modules have a dtype param, should we stray from that standard even though it is an activation function?

I created a constant init func because jax itself seemed to be lacking one, however I haven't received a response to the issue I posted in the jax repo requesting to add it so I'm fine with just using a lambda.
 - Other Modules have a dtype param to control the precision of their -intermediate- values, and a simple activation function like this doesn't have intermediates.  We don't require modules to surface a `dtype=` attribute - it's just convention for the core layers to do so to give users the ability to control the floating-point types of the "insides"

 - The "constant" functions you're looking for already exist: `jnp.full` and `jnp.full_like`
1. Makes sense thanks for clarifying that.
2. Thanks for pointing me jnp.full. I wasn't aware of that.

Shall I make a PR then?
Yeah if you'd like to make a PR we could add the above to `activations.py` I think (after all the passthrough function imports).  (but no pressure - if you don't have time we can add it soon ourselves.)
I'll try to take a first stab at it since it will be my first time contributing to flax.

Resolving patch:
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)


Test patch:
diff --git a/tests/linen/linen_activation_test.py b/tests/linen/linen_activation_test.py
new file mode 100644
--- /dev/null
+++ b/tests/linen/linen_activation_test.py
@@ -0,0 +1,42 @@
+# Copyright 2021 The Flax Authors.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+"""Tests for flax.nn.activation."""
+
+from absl.testing import absltest
+from absl.testing import parameterized
+
+from flax import linen as nn
+
+import jax
+from jax import random
+import jax.numpy as jnp
+
+
+# Parse absl flags test_srcdir and test_tmpdir.
+jax.config.parse_flags_with_absl()
+
+
+class ActivationTest(parameterized.TestCase):
+
+  def test_prelu(self):
+    rng = random.PRNGKey(0)
+    x = jnp.ones((4, 6, 5))
+    act = nn.PReLU()
+    y, _ = act.init_with_output(rng, x)
+    self.assertEqual(y.shape, x.shape)
+
+
+if __name__ == '__main__':
+  absltest.main()