code: 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)
  
imports: import pytest
  import numpy
  from tensorflow_addons.layers import Sparsemax
  from tensorflow_addons.utils import test_utils
label_desc: Use image node_id 1 for calling a global function with example usage: _np_sparsemax(z) and returns: np