def METHOD_NAME(inputs: tf.Tensor,
                   alpha: float = 1.5,
                   axis: int = -1,
                   n_iter: int = 50,
                   ensure_sum_one: bool = True) -> tf.Tensor:
  """Calculate the entmax probabilities."""
  @tf.custom_gradient
  def forward(inputs, alpha):
    with tf.name_scope("entmax_loss"):
      alpha_shape = inputs.get_shape().as_list()
      alpha_shape[axis] = 1
      alpha = tf.fill(alpha_shape, alpha)
      alpha = tf.cast(alpha, dtype=inputs.dtype)
      d = inputs.get_shape().as_list()[axis]
      alpha_m1 = alpha - 1.0
      inputs = inputs * alpha_m1
      max_val = tf.math.reduce_max(inputs, axis=axis, keepdims=True)
      tau_lo = max_val - tf.ones(
          alpha.get_shape().as_list(), dtype=inputs.dtype)
      tau_hi = max_val - tf.math.pow(
          tf.cast((1.0 / d), dtype=inputs.dtype), alpha_m1)
      f_lo = tf.math.reduce_sum(
          _calculate_probability(tf.math.subtract(inputs, tau_lo), alpha),
          axis) - 1.0
      dm = tau_hi - tau_lo
      for _ in range(n_iter):
        dm /= 2
        tau_m = tau_lo + dm
        p_m = _calculate_probability(inputs - tau_m, alpha)
        f_m = tf.math.reduce_sum(p_m, axis) - 1.0
        mask = tf.expand_dims(tf.math.greater(f_m * f_lo, 0), axis)
        tau_lo = tf.where(mask, tau_m, tau_lo)
      if ensure_sum_one:
        p_m /= tf.expand_dims(tf.math.reduce_sum(p_m, axis), axis)
    def grad_fn(d_outputs):
      with tf.name_scope("entmax_grad"):
        gppr = tf.where(p_m > 0, tf.math.pow(p_m, 2.0 - alpha),
                        tf.zeros_like(p_m))
        d_inputs = d_outputs * gppr
        q = tf.math.reduce_sum(d_inputs, axis) / tf.math.reduce_sum(gppr, axis)
        q = tf.expand_dims(q, axis)
        d_inputs -= q * gppr
        return d_inputs, d_inputs
    return p_m, grad_fn
  return forward(inputs, alpha)