Moved test out of run_all_in_graph_and_eager_mode in softshrink. (#1404)

See #1328

Author: gabrieldemarmiesse

tensorflow_addons/activations/softshrink_test.py

@@ -28,26 +28,28 @@
 from tensorflow_addons.utils import test_utils
 
 
-@test_utils.run_all_in_graph_and_eager_modes
-class SoftshrinkTest(tf.test.TestCase, parameterized.TestCase):
-    def test_invalid(self):
-        with self.assertRaisesOpError("lower must be less than or equal to upper."):
-            y = _softshrink_custom_op(tf.ones(shape=(1, 2, 3)), lower=2.0, upper=-2.0)
-            self.evaluate(y)
-
-    @parameterized.named_parameters(
-        ("float16", np.float16), ("float32", np.float32), ("float64", np.float64)
+def test_invalid():
+    with pytest.raises(
+        tf.errors.OpError, match="lower must be less than or equal to upper."
+    ):
+        y = _softshrink_custom_op(tf.ones(shape=(1, 2, 3)), lower=2.0, upper=-2.0)
+        y.numpy()
+
+
+@pytest.mark.parametrize("dtype", [np.float16, np.float32, np.float64])
+def test_softshrink(dtype):
+    x = tf.constant([-2.0, -1.0, 0.0, 1.0, 2.0], dtype=dtype)
+    expected_result = tf.constant([-1.5, -0.5, 0.0, 0.5, 1.5], dtype=dtype)
+    test_utils.assert_allclose_according_to_type(softshrink(x), expected_result)
+
+    expected_result = tf.constant([-1.0, 0.0, 0.0, 0.0, 1.0], dtype=dtype)
+    test_utils.assert_allclose_according_to_type(
+        softshrink(x, lower=-1.0, upper=1.0), expected_result
     )
-    def test_softshrink(self, dtype):
-        x = tf.constant([-2.0, -1.0, 0.0, 1.0, 2.0], dtype=dtype)
-        expected_result = tf.constant([-1.5, -0.5, 0.0, 0.5, 1.5], dtype=dtype)
-        self.assertAllCloseAccordingToType(softshrink(x), expected_result)
 
-        expected_result = tf.constant([-1.0, 0.0, 0.0, 0.0, 1.0], dtype=dtype)
-        self.assertAllCloseAccordingToType(
-            softshrink(x, lower=-1.0, upper=1.0), expected_result
-        )
 
+@test_utils.run_all_in_graph_and_eager_modes
+class SoftshrinkTest(tf.test.TestCase, parameterized.TestCase):
     @parameterized.named_parameters(("float32", np.float32), ("float64", np.float64))
     def test_theoretical_gradients(self, dtype):
         # Only test theoretical gradients for float32 and float64