//  [python]
def _paragraph_grouping_loss(
    loss_dict: Dict[str, Any],
    labels: Dict[str, Any],
    outputs: Dict[str, Any],
    tau: float = gin.REQUIRED,
    loss_mode="vanilla",
    fl_alpha: float = 0.25,
    fl_gamma: float = 2.,
):
  """Instance discrimination loss.

  This method adds the para discrimination loss term to loss_dict directly.

  Args:
    loss_dict: A dictionary for the loss. The values are loss scalars.
    labels: The label dictionary.
    outputs: The output dictionary.
    tau: The temperature term in the loss
    loss_mode: The type of loss.
    fl_alpha: alpha value in focal loss
    fl_gamma: gamma value in focal loss
  """
  if "paragraph_labels" not in labels:
    loss_dict["loss_para"] = 0.
    return
  # step 1:
  # obtain the paragraph labels for each prediction
  # (batch, pred, gt)
  matched_matrix = outputs["instance_output"]["matched_mask"]  # B, N, N
  para_label_gt = labels["paragraph_labels"]["paragraph_ids"]  # B, N
  has_para_label_gt = (
      labels["paragraph_labels"]["has_para_ids"][:, tf.newaxis, tf.newaxis])
  # '0' means no paragraph labels
  pred_label_gt = tf.einsum("bij,bj->bi", matched_matrix,
                            tf.cast(para_label_gt + 1, tf.float32))
  pred_label_gt_pad_col = tf.expand_dims(pred_label_gt, -1)  # b,n,1
  pred_label_gt_pad_row = tf.expand_dims(pred_label_gt, 1)  # b,1,n
  gt_affinity = tf.cast(
      tf.equal(pred_label_gt_pad_col, pred_label_gt_pad_row), tf.float32)
  gt_affinity_mask = (
      has_para_label_gt * pred_label_gt_pad_col * pred_label_gt_pad_row)
  gt_affinity_mask = tf.cast(tf.not_equal(gt_affinity_mask, 0.), tf.float32)

  # step 2:
  # get affinity matrix
  affinity = outputs["para_affinity"]

  # step 3:
  # compute loss
  loss_fn = tf_keras.losses.BinaryCrossentropy(
      from_logits=True,
      label_smoothing=0,
      axis=-1,
      reduction=tf_keras.losses.Reduction.NONE,
      name="para_dist")
  affinity = tf.reshape(affinity, (-1, 1))  # (b*n*n, 1)
  gt_affinity = tf.reshape(gt_affinity, (-1, 1))  # (b*n*n, 1)
  gt_affinity_mask = tf.reshape(gt_affinity_mask, (-1,))  # (b*n*n,)
  pointwise_loss = loss_fn(gt_affinity, affinity / tau)  # (b*n*n,)

  if loss_mode == "vanilla":
    loss = (
        tf.reduce_sum(pointwise_loss * gt_affinity_mask) /
        (tf.reduce_sum(gt_affinity_mask) + EPSILON))
  elif loss_mode == "balanced":
    # pos
    pos_mask = gt_affinity_mask * gt_affinity[:, 0]
    pos_loss = (
        tf.reduce_sum(pointwise_loss * pos_mask) /
        (tf.reduce_sum(pos_mask) + EPSILON))
    # neg
    neg_mask = gt_affinity_mask * (1. - gt_affinity[:, 0])
    neg_loss = (
        tf.reduce_sum(pointwise_loss * neg_mask) /
        (tf.reduce_sum(neg_mask) + EPSILON))
    loss = 0.25 * pos_loss + 0.75 * neg_loss
  elif loss_mode == "focal":
    alpha_wt = fl_alpha * gt_affinity + (1. - fl_alpha) * (1. - gt_affinity)
    prob_pos = tf.math.sigmoid(affinity / tau)
    pt = prob_pos * gt_affinity + (1. - prob_pos) * (1. - gt_affinity)
    fl_loss_pw = tf.stop_gradient(
        alpha_wt * tf.pow(1. - pt, fl_gamma))[:, 0] * pointwise_loss
    loss = (
        tf.reduce_sum(fl_loss_pw * gt_affinity_mask) /
        (tf.reduce_sum(gt_affinity_mask) + EPSILON))
  else:
    raise ValueError(f"Not supported loss mode: {loss_mode}")

  loss_dict["loss_para"] = loss