code: def sparsemax_loss(
      logits: TensorLike,
      sparsemax: TensorLike,
      labels: TensorLike,
      name: Optional[str] = None,
  ) -> tf.Tensor:
      """Sparsemax loss function [1].
  
      Computes the generalized multi-label classification loss for the sparsemax
      function. The implementation is a reformulation of the original loss
      function such that it uses the sparsemax probability output instead of the
      internal $ \tau $ variable. However, the output is identical to the original
      loss function.
  
      [1]: https://arxiv.org/abs/1602.02068
  
      Args:
        logits: A `Tensor`. Must be one of the following types: `float32`,
          `float64`.
        sparsemax: A `Tensor`. Must have the same type as `logits`.
        labels: A `Tensor`. Must have the same type as `logits`.
        name: A name for the operation (optional).
      Returns:
        A `Tensor`. Has the same type as `logits`.
      """
      logits = tf.convert_to_tensor(logits, name="logits")
      sparsemax = tf.convert_to_tensor(sparsemax, name="sparsemax")
      labels = tf.convert_to_tensor(labels, name="labels")
  
      # In the paper, they call the logits z.
      # A constant can be substracted from logits to make the algorithm
      # more numerically stable in theory. However, there are really no major
      # source numerical instability in this algorithm.
      z = logits
  
      # sum over support
      # Use a conditional where instead of a multiplication to support z = -inf.
      # If z = -inf, and there is no support (sparsemax = 0), a multiplication
      # would cause 0 * -inf = nan, which is not correct in this case.
      sum_s = tf.where(
          tf.math.logical_or(sparsemax > 0, tf.math.is_nan(sparsemax)),
          sparsemax * (z - 0.5 * sparsemax),
          tf.zeros_like(sparsemax),
      )
  
      # - z_k + ||q||^2
      q_part = labels * (0.5 * labels - z)
      # Fix the case where labels = 0 and z = -inf, where q_part would
      # otherwise be 0 * -inf = nan. But since the lables = 0, no cost for
      # z = -inf should be consideredself.
      # The code below also coveres the case where z = inf. Howeverm in this
      # caose the sparsemax will be nan, which means the sum_s will also be nan,
      # therefor this case doesn't need addtional special treatment.
      q_part_safe = tf.where(
          tf.math.logical_and(
              tf.math.equal(labels, 0), tf.math.is_inf(z)
          ),
          tf.zeros_like(z),
          q_part,
      )
  
      return tf.math.reduce_sum(sum_s + q_part_safe, axis=1)
  
file_path: addons/tensorflow_addons/losses/sparsemax_loss.py
imports: import tensorflow
  from tensorflow_addons.activations.sparsemax import sparsemax
  from tensorflow_addons.utils.types import TensorLike
  from typeguard import typechecked
  from typing import Optional
label_desc: Use image node_id 1 for calling a global function with example usage: sparsemax_loss(logits, sparsemax, labels, name) and returns: tf