text: # Copyright 2016 The TensorFlow Authors. All Rights Reserved.
  #
  # 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.
  # ==============================================================================
  
  import numpy as np
  import tensorflow as tf
  
  from tensorflow_addons.activations import sparsemax
  from tensorflow_addons.utils import test_utils
  
  test_obs = 17
  
  
  def _np_sparsemax(z):
      z = z - np.mean(z, axis=1)[:, np.newaxis]
  
      # sort z
      z_sorted = np.sort(z, axis=1)[:, ::-1]
  
      # calculate k(z)
      z_cumsum = np.cumsum(z_sorted, axis=1)
      k = np.arange(1, z.shape[1] + 1)
      z_check = 1 + k * z_sorted > z_cumsum
      # use argmax to get the index by row as .nonzero() doesn't
      # take an axis argument. np.argmax return the first index, but the last
      # index is required here, use np.flip to get the last index and
      # `z.shape[axis]` to compensate for np.flip afterwards.
      k_z = z.shape[1] - np.argmax(z_check[:, ::-1], axis=1)
  
      # calculate tau(z)
      tau_sum = z_cumsum[np.arange(0, z.shape[0]), k_z - 1]
      tau_z = ((tau_sum - 1) / k_z).reshape(-1, 1)
  
      # calculate p
      return np.maximum(0, z - tau_z)
  
  
  @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
  
      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))
  
          tf_sparsemax_op, tf_sparsemax_out = self._tf_sparsemax(z, dtype)
          np_sparsemax = _np_sparsemax(z).astype(dtype)
  
          self.assertAllCloseAccordingToType(np_sparsemax, tf_sparsemax_out)
          self.assertShapeEqual(np_sparsemax, tf_sparsemax_op)
  
      def test_sparsemax_against_numpy_axis(self, dtype=None):
          """check sparsemax kernel against numpy."""
          random = np.random.RandomState(1)
  
          z = random.uniform(low=-3, high=3, size=(test_obs, 10))
  
          tf_sparsemax_op, tf_sparsemax_out = self._tf_sparsemax(z, dtype, axis=0)
          np_sparsemax = np.transpose(_np_sparsemax(np.transpose(z))).astype(dtype)
  
          self.assertAllCloseAccordingToType(np_sparsemax, tf_sparsemax_out)
          self.assertShapeEqual(np_sparsemax, tf_sparsemax_op)
  
      def test_sparsemax_against_numpy_low_rank(self, dtype=None):
          """check sparsemax kernel against numpy."""
          random = np.random.RandomState(1)
  
          z = random.uniform(low=-3, high=3, size=(10))
  
          tf_sparsemax_op, tf_sparsemax_out = self._tf_sparsemax(z, dtype)
          np_sparsemax = np.reshape(_np_sparsemax(np.reshape(z, [1, 10])), [10]).astype(
              dtype
          )
  
          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)
  
          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)
  
          self.assertAllCloseAccordingToType(np_sparsemax, tf_sparsemax_out)
          self.assertShapeEqual(np_sparsemax, tf_sparsemax_op)
  
      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(
              [
                  [np.nan, np.nan, np.nan],
                  [np.nan, np.nan, np.nan],
                  [np.nan, np.nan, np.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(
              [
                  [np.nan, np.nan, np.nan],
                  [np.nan, np.nan, np.nan],
                  [np.nan, np.nan, np.nan],
              ],
              tf_sparsemax_pos,
          )
  
          _, tf_sparsemax_mix = self._tf_sparsemax(z_mix, dtype)
          self.assertAllEqual(
              [
                  [np.nan, np.nan, np.nan],
                  [np.nan, np.nan, np.nan],
                  [np.nan, np.nan, np.nan],
              ],
              tf_sparsemax_mix,
          )
  
      def test_sparsemax_of_zero(self, dtype=None):
          """check sparsemax proposition 1, part 1."""
          z = np.zeros((1, 10))
  
          tf_sparsemax_op, tf_sparsemax_out = self._tf_sparsemax(z, dtype)
          np_sparsemax = np.ones_like(z, dtype=dtype) / z.size
  
          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)
  
          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)
  
          # 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)
  
          self.assertAllCloseAccordingToType(p_expected, tf_sparsemax_out)
          self.assertShapeEqual(p_expected, tf_sparsemax_op)
  
      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)
  
          _, tf_sparsemax_zpc = self._tf_sparsemax(z + c, dtype)
  
          _, tf_sparsemax_z = self._tf_sparsemax(z, dtype)
  
          self.assertAllCloseAccordingToType(
              tf_sparsemax_zpc, tf_sparsemax_z, half_atol=5e-3
          )
  
      def test_permutation(self, dtype=None):
          """check sparsemax proposition 3."""
          random = np.random.RandomState(6)
  
          z = random.uniform(low=-3, high=3, size=(test_obs, 10))
          _, p = self._tf_sparsemax(z, dtype)
  
          for i in range(test_obs):
              per = random.permutation(10)
  
              tf_sparsemax_op, tf_sparsemax_out = self._tf_sparsemax(
                  z[i, per].reshape(1, -1), dtype
              )
              p_expected = p[i, per].reshape(1, -1)
  
              self.assertAllCloseAccordingToType(
                  p_expected, tf_sparsemax_out, half_atol=5e-3
              )
              self.assertShapeEqual(p_expected, tf_sparsemax_op)
  
      def test_diffrence(self, dtype=None):
          """check sparsemax proposition 4."""
          random = np.random.RandomState(7)
  
          z = random.uniform(low=-3, high=3, size=(test_obs, 10))
          _, p = self._tf_sparsemax(z, dtype)
  
          etol = {"float16": 1e-2, "float32": 1e-6, "float64": 1e-9}[dtype]
  
          for val in range(0, test_obs):
              for i in range(0, 10):
                  for j in range(0, 10):
                      # check condition, the obesite pair will be checked anyway
                      if z[val, i] > z[val, j]:
                          continue
  
                      self.assertTrue(
                          0 <= p[val, j] - p[val, i] <= z[val, j] - z[val, i] + etol,
                          "0 <= %.10f <= %.10f"
                          % (p[val, j] - p[val, i], z[val, j] - z[val, i] + etol),
                      )
  
      def test_two_dimentional(self, dtype=None):
          """check two dimentation sparsemax case."""
          t = np.linspace(-2, 2, test_obs, dtype=dtype)
          z = np.vstack([t, np.zeros(test_obs, dtype=dtype)]).T
  
          tf_sparsemax_op, tf_sparsemax_out = self._tf_sparsemax(z, dtype)
  
          p0_expected = np.select([t < -1, t <= 1, t > 1], [0, (t + 1) / 2, 1])
  
          self.assertAllCloseAccordingToType(p0_expected, tf_sparsemax_out[:, 0])
          self.assertAllCloseAccordingToType(1 - p0_expected, tf_sparsemax_out[:, 1])
          self.assertShapeEqual(z, tf_sparsemax_op)
  
      def test_gradient_against_estimate(self, dtype=None):
          """check sparsemax Rop, against estimated Rop."""
          random = np.random.RandomState(9)
  
          # sparsemax is not a smooth function so gradient estimation is only
          # possible for float64.
          if dtype != "float64":
              return
  
          z = random.uniform(low=-1, high=1, size=(test_obs, 10)).astype(dtype)
  
          (jacob_sym,), (jacob_num,) = tf.test.compute_gradient(
              lambda logits: sparsemax(logits), [z], delta=1e-6
          )
          self.assertAllCloseAccordingToType(jacob_sym, jacob_num)
  
  
  if __name__ == "__main__":
      tf.test.main()
  