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

add linen activation tests