Removed run_all_in_graph_and_eager_mode in sparsemax. (#1689)

* Removed run_all_in_graph_and_eager_mode in sparsemax.

* Removed old code.

* Removed useless astype.

Author: gabrieldemarmiesse

tensorflow_addons/activations/tests/sparsemax_test.py

@@ -79,140 +79,146 @@ def test_sparsemax_against_numpy_low_rank(dtype):
     assert np_sparsemax.shape == tf_sparsemax_out.shape
 
 
-@test_utils.run_all_with_types(["float32", "float64"])
-@test_utils.run_all_in_graph_and_eager_modes
-class SparsemaxTest(tf.test.TestCase):
-    def _tf_sparsemax(self, z, dtype, **kwargs):
-        tf_sparsemax_op = sparsemax(z.astype(dtype), **kwargs)
-        tf_sparsemax_out = self.evaluate(tf_sparsemax_op)
-
-        return tf_sparsemax_op, tf_sparsemax_out
+@pytest.mark.parametrize("dtype", ["float32", "float64"])
+def test_sparsemax_against_numpy(dtype):
+    """check sparsemax kernel against numpy."""
+    random = np.random.RandomState(1)
 
-    def test_sparsemax_against_numpy(self, dtype=None):
-        """check sparsemax kernel against numpy."""
-        random = np.random.RandomState(1)
+    z = random.uniform(low=-3, high=3, size=(test_obs, 10))
 
-        z = random.uniform(low=-3, high=3, size=(test_obs, 10))
+    tf_sparsemax_out = sparsemax(z.astype(dtype))
+    np_sparsemax = _np_sparsemax(z).astype(dtype)
 
-        tf_sparsemax_op, tf_sparsemax_out = self._tf_sparsemax(z, dtype)
-        np_sparsemax = _np_sparsemax(z).astype(dtype)
+    test_utils.assert_allclose_according_to_type(np_sparsemax, tf_sparsemax_out)
 
-        self.assertAllCloseAccordingToType(np_sparsemax, tf_sparsemax_out)
-        self.assertShapeEqual(np_sparsemax, tf_sparsemax_op)
 
-    def test_sparsemax_against_numpy_high_rank(self, dtype=None):
-        """check sparsemax kernel against numpy."""
-        random = np.random.RandomState(1)
+@pytest.mark.parametrize("dtype", ["float32", "float64"])
+def test_sparsemax_against_numpy_high_rank(dtype):
+    """check sparsemax kernel against numpy."""
+    random = np.random.RandomState(1)
 
-        z = random.uniform(low=-3, high=3, size=(test_obs, test_obs, 10))
+    z = random.uniform(low=-3, high=3, size=(test_obs, test_obs, 10))
 
-        tf_sparsemax_op, tf_sparsemax_out = self._tf_sparsemax(z, dtype)
-        np_sparsemax = np.reshape(
-            _np_sparsemax(np.reshape(z, [test_obs * test_obs, 10])),
-            [test_obs, test_obs, 10],
-        ).astype(dtype)
+    tf_sparsemax_out = sparsemax(z.astype(dtype))
+    np_sparsemax = np.reshape(
+        _np_sparsemax(np.reshape(z, [test_obs * test_obs, 10])),
+        [test_obs, test_obs, 10],
+    ).astype(dtype)
 
-        self.assertAllCloseAccordingToType(np_sparsemax, tf_sparsemax_out)
-        self.assertShapeEqual(np_sparsemax, tf_sparsemax_op)
+    test_utils.assert_allclose_according_to_type(np_sparsemax, tf_sparsemax_out)
 
-    def test_sparsemax_of_nan(self, dtype=None):
-        """check sparsemax transfers nan."""
-        z_nan = np.asarray(
-            [[0, np.nan, 0], [0, np.nan, np.nan], [np.nan, np.nan, np.nan],]
-        ).astype(dtype)
 
-        _, tf_sparsemax_nan = self._tf_sparsemax(z_nan, dtype)
-        self.assertAllEqual(
+@pytest.mark.parametrize("dtype", ["float32", "float64"])
+def test_sparsemax_of_nan(dtype):
+    """check sparsemax transfers nan."""
+    z_nan = np.asarray(
+        [[0, np.nan, 0], [0, np.nan, np.nan], [np.nan, np.nan, np.nan],]
+    ).astype(dtype)
+
+    tf_sparsemax_nan = sparsemax(z_nan)
+    np.testing.assert_equal(
+        np.array(
             [
                 [np.nan, np.nan, np.nan],
                 [np.nan, np.nan, np.nan],
                 [np.nan, np.nan, np.nan],
-            ],
-            tf_sparsemax_nan,
-        )
+            ]
+        ),
+        tf_sparsemax_nan,
+    )
 
-    def test_sparsemax_of_inf(self, dtype=None):
-        """check sparsemax is infinity safe."""
-        z_neg = np.asarray(
-            [[0, -np.inf, 0], [0, -np.inf, -np.inf], [-np.inf, -np.inf, -np.inf],]
-        ).astype(dtype)
-        z_pos = np.asarray(
-            [[0, np.inf, 0], [0, np.inf, np.inf], [np.inf, np.inf, np.inf]]
-        ).astype(dtype)
-        z_mix = np.asarray(
-            [[0, np.inf, 0], [0, np.inf, -np.inf], [-np.inf, np.inf, -np.inf]]
-        ).astype(dtype)
-
-        _, tf_sparsemax_neg = self._tf_sparsemax(z_neg, dtype)
-        self.assertAllEqual(
-            [[0.5, 0, 0.5], [1, 0, 0], [np.nan, np.nan, np.nan]], tf_sparsemax_neg
-        )
 
-        _, tf_sparsemax_pos = self._tf_sparsemax(z_pos, dtype)
-        self.assertAllEqual(
+@pytest.mark.parametrize("dtype", ["float32", "float64"])
+def test_sparsemax_of_inf(dtype):
+    """check sparsemax is infinity safe."""
+    z_neg = np.asarray(
+        [[0, -np.inf, 0], [0, -np.inf, -np.inf], [-np.inf, -np.inf, -np.inf],]
+    ).astype(dtype)
+    z_pos = np.asarray(
+        [[0, np.inf, 0], [0, np.inf, np.inf], [np.inf, np.inf, np.inf]]
+    ).astype(dtype)
+    z_mix = np.asarray(
+        [[0, np.inf, 0], [0, np.inf, -np.inf], [-np.inf, np.inf, -np.inf]]
+    ).astype(dtype)
+
+    tf_sparsemax_neg = sparsemax(z_neg)
+    np.testing.assert_equal(
+        np.array([[0.5, 0, 0.5], [1, 0, 0], [np.nan, np.nan, np.nan]]), tf_sparsemax_neg
+    )
+
+    tf_sparsemax_pos = sparsemax(z_pos)
+    np.testing.assert_equal(
+        np.array(
             [
                 [np.nan, np.nan, np.nan],
                 [np.nan, np.nan, np.nan],
                 [np.nan, np.nan, np.nan],
-            ],
-            tf_sparsemax_pos,
-        )
+            ]
+        ),
+        tf_sparsemax_pos,
+    )
 
-        _, tf_sparsemax_mix = self._tf_sparsemax(z_mix, dtype)
-        self.assertAllEqual(
+    tf_sparsemax_mix = sparsemax(z_mix)
+    np.testing.assert_equal(
+        np.array(
             [
                 [np.nan, np.nan, np.nan],
                 [np.nan, np.nan, np.nan],
                 [np.nan, np.nan, np.nan],
-            ],
-            tf_sparsemax_mix,
-        )
+            ]
+        ),
+        tf_sparsemax_mix,
+    )
+
+
+@pytest.mark.parametrize("dtype", ["float32", "float64"])
+def test_sparsemax_of_zero(dtype):
+    """check sparsemax proposition 1, part 1."""
+    z = np.zeros((1, 10))
 
-    def test_sparsemax_of_zero(self, dtype=None):
-        """check sparsemax proposition 1, part 1."""
-        z = np.zeros((1, 10))
+    tf_sparsemax_out = sparsemax(z.astype(dtype))
+    np_sparsemax = np.ones_like(z, dtype=dtype) / z.size
 
-        tf_sparsemax_op, tf_sparsemax_out = self._tf_sparsemax(z, dtype)
-        np_sparsemax = np.ones_like(z, dtype=dtype) / z.size
+    test_utils.assert_allclose_according_to_type(np_sparsemax, tf_sparsemax_out)
 
-        self.assertAllCloseAccordingToType(np_sparsemax, tf_sparsemax_out)
-        self.assertShapeEqual(np_sparsemax, tf_sparsemax_op)
 
-    def test_sparsemax_of_to_inf(self, dtype=None):
-        """check sparsemax proposition 1, part 2."""
-        random = np.random.RandomState(4)
+@pytest.mark.parametrize("dtype", ["float32", "float64"])
+def test_sparsemax_of_to_inf(dtype):
+    """check sparsemax proposition 1, part 2."""
+    random = np.random.RandomState(4)
 
-        z = random.uniform(low=-3, high=3, size=(test_obs, 10))
+    z = random.uniform(low=-3, high=3, size=(test_obs, 10))
 
-        # assume |A(z)| = 1, as z is continues random
-        z_sort_arg = np.argsort(z, axis=1)[:, ::-1]
-        z_sort = np.sort(z, axis=-1)[:, ::-1]
-        gamma_z = z_sort[:, 0] - z_sort[:, 1]
-        epsilon = (0.99 * gamma_z * 1).reshape(-1, 1)
+    # assume |A(z)| = 1, as z is continues random
+    z_sort_arg = np.argsort(z, axis=1)[:, ::-1]
+    z_sort = np.sort(z, axis=-1)[:, ::-1]
+    gamma_z = z_sort[:, 0] - z_sort[:, 1]
+    epsilon = (0.99 * gamma_z * 1).reshape(-1, 1)
 
-        # construct the expected 1_A(z) array
-        p_expected = np.zeros((test_obs, 10), dtype=dtype)
-        p_expected[np.arange(0, test_obs), z_sort_arg[:, 0]] = 1
+    # construct the expected 1_A(z) array
+    p_expected = np.zeros((test_obs, 10), dtype=dtype)
+    p_expected[np.arange(0, test_obs), z_sort_arg[:, 0]] = 1
 
-        tf_sparsemax_op, tf_sparsemax_out = self._tf_sparsemax((1 / epsilon) * z, dtype)
+    tf_sparsemax_out = sparsemax(((1 / epsilon) * z).astype(dtype))
 
-        self.assertAllCloseAccordingToType(p_expected, tf_sparsemax_out)
-        self.assertShapeEqual(p_expected, tf_sparsemax_op)
+    test_utils.assert_allclose_according_to_type(p_expected, tf_sparsemax_out)
 
-    def test_constant_add(self, dtype=None):
-        """check sparsemax proposition 2."""
-        random = np.random.RandomState(5)
 
-        z = random.uniform(low=-3, high=3, size=(test_obs, 10)).astype(dtype)
-        c = random.uniform(low=-3, high=3, size=(test_obs, 1)).astype(dtype)
+@pytest.mark.parametrize("dtype", ["float32", "float64"])
+def test_constant_add(dtype):
+    """check sparsemax proposition 2."""
+    random = np.random.RandomState(5)
 
-        _, tf_sparsemax_zpc = self._tf_sparsemax(z + c, dtype)
+    z = random.uniform(low=-3, high=3, size=(test_obs, 10)).astype(dtype)
+    c = random.uniform(low=-3, high=3, size=(test_obs, 1)).astype(dtype)
 
-        _, tf_sparsemax_z = self._tf_sparsemax(z, dtype)
+    tf_sparsemax_zpc = sparsemax((z + c))
 
-        self.assertAllCloseAccordingToType(
-            tf_sparsemax_zpc, tf_sparsemax_z, half_atol=5e-3
-        )
+    tf_sparsemax_z = sparsemax(z)
+
+    test_utils.assert_allclose_according_to_type(
+        tf_sparsemax_zpc, tf_sparsemax_z, half_atol=5e-3
+    )
 
 
 @pytest.mark.parametrize("dtype", ["float32", "float64"])