What would be equivalent of pytorch's torch.nn.CosineEmbeddingLoss in tensorflow?

A TensorFlow version of CosineEmbeddingLoss: 

import tensorflow as tf

def CosineEmbeddingLoss(margin=0.):
  def _cosine_similarity(x1, x2):
    """Cosine similarity between two batches of vectors."""
    return tf.reduce_sum(x1 * x2, axis=-1) / (
        tf.norm(x1, axis=-1) * tf.norm(x2, axis=-1))
  def _cosine_embedding_loss_fn(input_one, input_two, target):
    similarity = _cosine_similarity(input_one, input_two)
    return tf.reduce_mean(tf.where(
        tf.equal(target, 1),
        1. - similarity,
        tf.maximum(tf.zeros_like(similarity), similarity - margin)))
  return _cosine_embedding_loss_fn