code: def sparsemax(logits: types.TensorLike, axis: int = -1) -> tf.Tensor:
      r"""Sparsemax activation function.
  
      For each batch $i$, and class $j$,
      compute sparsemax activation function:
  
      $$
      \mathrm{sparsemax}(x)[i, j] = \max(\mathrm{logits}[i, j] - \tau(\mathrm{logits}[i, :]), 0).
      $$
  
      See [From Softmax to Sparsemax: A Sparse Model of Attention and Multi-Label Classification](https://arxiv.org/abs/1602.02068).
  
      Usage:
  
      >>> x = tf.constant([[-1.0, 0.0, 1.0], [-5.0, 1.0, 2.0]])
      >>> tfa.activations.sparsemax(x)
      <tf.Tensor: shape=(2, 3), dtype=float32, numpy=
      array([[0., 0., 1.],
             [0., 0., 1.]], dtype=float32)>
  
      Args:
          logits: A `Tensor`.
          axis: `int`, axis along which the sparsemax operation is applied.
      Returns:
          A `Tensor`, output of sparsemax transformation. Has the same type and
          shape as `logits`.
      Raises:
          ValueError: In case `dim(logits) == 1`.
      """
      logits = tf.convert_to_tensor(logits, name="logits")
  
      # We need its original shape for shape inference.
      shape = logits.get_shape()
      rank = shape.rank
      is_last_axis = (axis == -1) or (axis == rank - 1)
  
      if is_last_axis:
          output = _compute_2d_sparsemax(logits)
          output.set_shape(shape)
          return output
  
      # If dim is not the last dimension, we have to do a transpose so that we can
      # still perform softmax on its last dimension.
  
      # Swap logits' dimension of dim and its last dimension.
      rank_op = tf.rank(logits)
      axis_norm = axis % rank
      logits = _swap_axis(
          logits, axis_norm, tf.math.subtract(rank_op, 1)
      )
  
      # Do the actual softmax on its last dimension.
      output = _compute_2d_sparsemax(logits)
      output = _swap_axis(
          output, axis_norm, tf.math.subtract(rank_op, 1)
      )
  
      # Make shape inference work since transpose may erase its static shape.
      output.set_shape(shape)
      return output
  
imports: import tensorflow
  from tensorflow_addons.utils import types
label_desc: Use image node_id 1 for calling a global function with example usage: sparsemax(logits, axis) and returns: output, output