import tensorflow as tf

  def compute_triplet_loss(positive_distances, negative_distances):
    losses = tf.nn.relu(positive_distances + margin - negative_distances)
    losses = tf.where(
        tf.stop_gradient(losses < losses.dtype.max), losses,
        tf.zeros_like(losses))
    num_nonzero_losses = tf.math.count_nonzero(losses)
    loss = tf.math.reduce_mean(losses)
    return loss, num_nonzero_losses
