import tensorflow as tf

  loss = tf.math.reduce_mean(
      distance_utils.compute_gaussian_kl_divergence(
          means, stddevs, rhs_means=prior_mean, rhs_stddevs=prior_stddev))
  weighted_loss = loss_weight * loss
  summaries = {
      'regularization_loss/KL/PriorMean/Mean':
          tf.math.reduce_mean(tf.constant(prior_mean)),
      'regularization_loss/KL/PriorVar/Mean':
          tf.math.reduce_mean(tf.constant(prior_stddev)**2),
      'regularization_loss/KL/Loss/Original':
          loss,
      'regularization_loss/KL/Loss/Weighted':
          weighted_loss,
      'regularization_loss/KL/Loss/Weight':
          tf.constant(loss_weight),
  }
  return weighted_loss, summaries


def compute_positive_pairwise_loss(anchor_embeddings,
                                   positive_embeddings,
                                   loss_weight,
                                   distance_fn=functools.partial(
                                       distance_utils.compute_l2_distances,
                                       squared=True)):
  """Computes anchor/positive pairwise (squared L2) loss.

  Args:
