import tensorflow as tf

def nt_xent(z1, z2, batch_size, temperature, zdim):
    """Implements normalized temperature-scaled cross-entropy loss

    Args:
        z1: normalized latent representation of first set of augmented images [N, D]
        z2: normalized latent representation of second set of augmented images [N, D]
        batch_size: number of images in batch
        temperature: temperature for softmax. set in config
        zdim: dimension for latent representation set in config
    Returns:
        loss: contrastive loss averaged over batch (2*N samples)
    """

    # reshape so that the order is z1_1,z2_1,z1_2,z2_2,z1_3,z2_3
    z = tf.concat([z1,z2], axis=0)
    z_ = tf.reshape(tf.transpose(tf.reshape(z, [2, batch_size, zdim]), [1,0,2]), [batch_size*2, -1])

    # compute cosine similarity
    # a has order [z1_1*batch_size*2, z1_2*batch_size*2, ...]
    # b has order [z1_1, z1_2, z3_1 ...]
    a = tf.reshape(tf.transpose(tf.tile(tf.reshape(z_, [1, batch_size*2, zdim]), [batch_size*2 ,1, 1]), [1, 0, 2]), [batch_size*2*batch_size*2, zdim])
    b = tf.tile(z_, [batch_size*2, 1])
    sim = cosine_similarity(a, b)
    sim = tf.expand_dims(sim, axis=1)/temperature
    sim = tf.reshape(sim, [batch_size*2, batch_size*2])
    sim = tf.math.exp(sim-tf.reduce_max(sim))

    pos_indices = tf.concat([tf.range(1, (2*batch_size)**2, (batch_size*4)+2), tf.range(batch_size*2, (2*batch_size)**2, (batch_size*4)+2)], axis=0)
    pos_indices = tf.expand_dims(pos_indices, axis=1)
    pos_mask = tf.zeros(((2*batch_size)**2, 1), dtype=tf.int32)
    pos_mask = tf.tensor_scatter_nd_add(pos_mask, pos_indices, tf.ones((batch_size*2, 1), dtype=tf.int32))
    pos_mask = tf.reshape(pos_mask, [batch_size*2, batch_size*2])
    neg_mask = tf.ones((batch_size*2, batch_size*2), dtype=tf.int32) - tf.eye(batch_size*2, dtype=tf.int32)

    # similarity between z11-z12, z12-z11, z21-22, z22-z21 etc. 
    pos_sim = tf.reduce_sum(sim*tf.cast(pos_mask, tf.float32), axis=1) 

    # negative similarity consists of all similarities except i=j
    neg_sim = tf.reduce_sum(sim*tf.cast(neg_mask, tf.float32), axis=1)
    loss = -tf.reduce_mean(tf.math.log(tf.clip_by_value(pos_sim/neg_sim, 1e-10, 1.0)))

    return loss

def cosine_similarity(a, b):
    """Computes the cosine similarity between vectors a and b"""

    numerator = tf.reduce_sum(tf.multiply(a, b), axis=1)
    denominator = tf.multiply(tf.norm(a, axis=1), tf.norm(b, axis=1))
    cos_similarity = numerator/denominator
    return cos_similarity


fca1ac97df2b7d7d0169848221c44c335c0c136c
<mask token>

<mask token>


def nt_xent(z1, z2, batch_size, temperature, zdim):
    """Implements normalized temperature-scaled cross-entropy loss

    Args:
        z1: normalized latent representation of first set of augmented images [N, D]
        z2: normalized latent representation of second set of augmented images [N, D]
        batch_size: number of images in batch
        temperature: temperature for softmax. set in config
        zdim: dimension for latent representation set in config
    Returns:
        loss: contrastive loss averaged over batch (2*N samples)
    """
    z = tf.concat([z1, z2], axis=0)
    z_ = tf.reshape(tf.transpose(tf.reshape(z, [2, batch_size, zdim]), [1, 
        0, 2]), [batch_size * 2, -1])
    a = tf.reshape(tf.transpose(tf.tile(tf.reshape(z_, [1, batch_size * 2,
        zdim]), [batch_size * 2, 1, 1]), [1, 0, 2]), [batch_size * 2 *
        batch_size * 2, zdim])
    b = tf.tile(z_, [batch_size * 2, 1])
    sim = cosine_similarity(a, b)
    sim = tf.expand_dims(sim, axis=1) / temperature
    sim = tf.reshape(sim, [batch_size * 2, batch_size * 2])
    sim = tf.math.exp(sim - tf.reduce_max(sim))
    pos_indices = tf.concat([tf.range(1, (2 * batch_size) ** 2, batch_size *
        4 + 2), tf.range(batch_size * 2, (2 * batch_size) ** 2, batch_size *
        4 + 2)], axis=0)
    pos_indices = tf.expand_dims(pos_indices, axis=1)
    pos_mask = tf.zeros(((2 * batch_size) ** 2, 1), dtype=tf.int32)
    pos_mask = tf.tensor_scatter_nd_add(pos_mask, pos_indices, tf.ones((
        batch_size * 2, 1), dtype=tf.int32))
    pos_mask = tf.reshape(pos_mask, [batch_size * 2, batch_size * 2])
    neg_mask = tf.ones((batch_size * 2, batch_size * 2), dtype=tf.int32
        ) - tf.eye(batch_size * 2, dtype=tf.int32)
    pos_sim = tf.reduce_sum(sim * tf.cast(pos_mask, tf.float32), axis=1)
    neg_sim = tf.reduce_sum(sim * tf.cast(neg_mask, tf.float32), axis=1)
    loss = -tf.reduce_mean(tf.math.log(tf.clip_by_value(pos_sim / neg_sim, 
        1e-10, 1.0)))
    return loss


<mask token>

<mask token>


def nt_xent(z1, z2, batch_size, temperature, zdim):
    """Implements normalized temperature-scaled cross-entropy loss

    Args:
        z1: normalized latent representation of first set of augmented images [N, D]
        z2: normalized latent representation of second set of augmented images [N, D]
        batch_size: number of images in batch
        temperature: temperature for softmax. set in config
        zdim: dimension for latent representation set in config
    Returns:
        loss: contrastive loss averaged over batch (2*N samples)
    """
    z = tf.concat([z1, z2], axis=0)
    z_ = tf.reshape(tf.transpose(tf.reshape(z, [2, batch_size, zdim]), [1, 
        0, 2]), [batch_size * 2, -1])
    a = tf.reshape(tf.transpose(tf.tile(tf.reshape(z_, [1, batch_size * 2,
        zdim]), [batch_size * 2, 1, 1]), [1, 0, 2]), [batch_size * 2 *
        batch_size * 2, zdim])
    b = tf.tile(z_, [batch_size * 2, 1])
    sim = cosine_similarity(a, b)
    sim = tf.expand_dims(sim, axis=1) / temperature
    sim = tf.reshape(sim, [batch_size * 2, batch_size * 2])
    sim = tf.math.exp(sim - tf.reduce_max(sim))
    pos_indices = tf.concat([tf.range(1, (2 * batch_size) ** 2, batch_size *
        4 + 2), tf.range(batch_size * 2, (2 * batch_size) ** 2, batch_size *
        4 + 2)], axis=0)
    pos_indices = tf.expand_dims(pos_indices, axis=1)
    pos_mask = tf.zeros(((2 * batch_size) ** 2, 1), dtype=tf.int32)
    pos_mask = tf.tensor_scatter_nd_add(pos_mask, pos_indices, tf.ones((
        batch_size * 2, 1), dtype=tf.int32))
    pos_mask = tf.reshape(pos_mask, [batch_size * 2, batch_size * 2])
    neg_mask = tf.ones((batch_size * 2, batch_size * 2), dtype=tf.int32
        ) - tf.eye(batch_size * 2, dtype=tf.int32)
    pos_sim = tf.reduce_sum(sim * tf.cast(pos_mask, tf.float32), axis=1)
    neg_sim = tf.reduce_sum(sim * tf.cast(neg_mask, tf.float32), axis=1)
    loss = -tf.reduce_mean(tf.math.log(tf.clip_by_value(pos_sim / neg_sim, 
        1e-10, 1.0)))
    return loss


def cosine_similarity(a, b):
    """Computes the cosine similarity between vectors a and b"""
    numerator = tf.reduce_sum(tf.multiply(a, b), axis=1)
    denominator = tf.multiply(tf.norm(a, axis=1), tf.norm(b, axis=1))
    cos_similarity = numerator / denominator
    return cos_similarity

import tensorflow as tf


def nt_xent(z1, z2, batch_size, temperature, zdim):
    """Implements normalized temperature-scaled cross-entropy loss

    Args:
        z1: normalized latent representation of first set of augmented images [N, D]
        z2: normalized latent representation of second set of augmented images [N, D]
        batch_size: number of images in batch
        temperature: temperature for softmax. set in config
        zdim: dimension for latent representation set in config
    Returns:
        loss: contrastive loss averaged over batch (2*N samples)
    """
    z = tf.concat([z1, z2], axis=0)
    z_ = tf.reshape(tf.transpose(tf.reshape(z, [2, batch_size, zdim]), [1, 
        0, 2]), [batch_size * 2, -1])
    a = tf.reshape(tf.transpose(tf.tile(tf.reshape(z_, [1, batch_size * 2,
        zdim]), [batch_size * 2, 1, 1]), [1, 0, 2]), [batch_size * 2 *
        batch_size * 2, zdim])
    b = tf.tile(z_, [batch_size * 2, 1])
    sim = cosine_similarity(a, b)
    sim = tf.expand_dims(sim, axis=1) / temperature
    sim = tf.reshape(sim, [batch_size * 2, batch_size * 2])
    sim = tf.math.exp(sim - tf.reduce_max(sim))
    pos_indices = tf.concat([tf.range(1, (2 * batch_size) ** 2, batch_size *
        4 + 2), tf.range(batch_size * 2, (2 * batch_size) ** 2, batch_size *
        4 + 2)], axis=0)
    pos_indices = tf.expand_dims(pos_indices, axis=1)
    pos_mask = tf.zeros(((2 * batch_size) ** 2, 1), dtype=tf.int32)
    pos_mask = tf.tensor_scatter_nd_add(pos_mask, pos_indices, tf.ones((
        batch_size * 2, 1), dtype=tf.int32))
    pos_mask = tf.reshape(pos_mask, [batch_size * 2, batch_size * 2])
    neg_mask = tf.ones((batch_size * 2, batch_size * 2), dtype=tf.int32
        ) - tf.eye(batch_size * 2, dtype=tf.int32)
    pos_sim = tf.reduce_sum(sim * tf.cast(pos_mask, tf.float32), axis=1)
    neg_sim = tf.reduce_sum(sim * tf.cast(neg_mask, tf.float32), axis=1)
    loss = -tf.reduce_mean(tf.math.log(tf.clip_by_value(pos_sim / neg_sim, 
        1e-10, 1.0)))
    return loss


def cosine_similarity(a, b):
    """Computes the cosine similarity between vectors a and b"""
    numerator = tf.reduce_sum(tf.multiply(a, b), axis=1)
    denominator = tf.multiply(tf.norm(a, axis=1), tf.norm(b, axis=1))
    cos_similarity = numerator / denominator
    return cos_similarity

import tensorflow as tf

def nt_xent(z1, z2, batch_size, temperature, zdim):
    """Implements normalized temperature-scaled cross-entropy loss

    Args:
        z1: normalized latent representation of first set of augmented images [N, D]
        z2: normalized latent representation of second set of augmented images [N, D]
        batch_size: number of images in batch
        temperature: temperature for softmax. set in config
        zdim: dimension for latent representation set in config
    Returns:
        loss: contrastive loss averaged over batch (2*N samples)
    """

    # reshape so that the order is z1_1,z2_1,z1_2,z2_2,z1_3,z2_3
    z = tf.concat([z1,z2], axis=0)
    z_ = tf.reshape(tf.transpose(tf.reshape(z, [2, batch_size, zdim]), [1,0,2]), [batch_size*2, -1])

    # compute cosine similarity
    # a has order [z1_1*batch_size*2, z1_2*batch_size*2, ...]
    # b has order [z1_1, z1_2, z3_1 ...]
    a = tf.reshape(tf.transpose(tf.tile(tf.reshape(z_, [1, batch_size*2, zdim]), [batch_size*2 ,1, 1]), [1, 0, 2]), [batch_size*2*batch_size*2, zdim])
    b = tf.tile(z_, [batch_size*2, 1])
    sim = cosine_similarity(a, b)
    sim = tf.expand_dims(sim, axis=1)/temperature
    sim = tf.reshape(sim, [batch_size*2, batch_size*2])
    sim = tf.math.exp(sim-tf.reduce_max(sim))

    pos_indices = tf.concat([tf.range(1, (2*batch_size)**2, (batch_size*4)+2), tf.range(batch_size*2, (2*batch_size)**2, (batch_size*4)+2)], axis=0)
    pos_indices = tf.expand_dims(pos_indices, axis=1)
    pos_mask = tf.zeros(((2*batch_size)**2, 1), dtype=tf.int32)
    pos_mask = tf.tensor_scatter_nd_add(pos_mask, pos_indices, tf.ones((batch_size*2, 1), dtype=tf.int32))
    pos_mask = tf.reshape(pos_mask, [batch_size*2, batch_size*2])
    neg_mask = tf.ones((batch_size*2, batch_size*2), dtype=tf.int32) - tf.eye(batch_size*2, dtype=tf.int32)

    # similarity between z11-z12, z12-z11, z21-22, z22-z21 etc. 
    pos_sim = tf.reduce_sum(sim*tf.cast(pos_mask, tf.float32), axis=1) 

    # negative similarity consists of all similarities except i=j
    neg_sim = tf.reduce_sum(sim*tf.cast(neg_mask, tf.float32), axis=1)
    loss = -tf.reduce_mean(tf.math.log(tf.clip_by_value(pos_sim/neg_sim, 1e-10, 1.0)))

    return loss

def cosine_similarity(a, b):
    """Computes the cosine similarity between vectors a and b"""

    numerator = tf.reduce_sum(tf.multiply(a, b), axis=1)
    denominator = tf.multiply(tf.norm(a, axis=1), tf.norm(b, axis=1))
    cos_similarity = numerator/denominator
    return cos_similarity

