What would be equivalent of pytorch's torch.nn.CosineEmbeddingLoss in tensorflow?
CosineEmbeddingLoss in Pytorch is the perfect function I am looking for in tensorflow, but I can only find tf.losses.cosine_distance. Is there a way or code that writes 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


Running it alongside Torch's version:

import tensorflow as tf
import numpy
import torch
from torch.autograd import Variable


first_values = numpy.random.normal(size=[100, 3])
second_values = numpy.random.normal(size=[100, 3])
labels = numpy.random.randint(2, size=[100]) * 2 - 1
torch_result = torch.nn.CosineEmbeddingLoss(margin=0.5)(
    Variable(torch.FloatTensor(first_values)),
    Variable(torch.FloatTensor(second_values)),
    Variable(torch.IntTensor(labels))).data.numpy()
with tf.Graph().as_default():
  with tf.Session():
    tf_result = CosineEmbeddingLoss(margin=0.5)(
        first_values, second_values, labels).eval()
print(torch_result, tf_result)


Seems to match to within reasonable precision:

array([ 0.35702518], dtype=float32)  0.35702516587462357

