import tensorflow as tf

    sims_logits = tf.matmul(z_projs, z_aug_projs, transpose_b=True)
    logits_max = tf.reduce_max(sims_logits,1)
    sims_logits = sims_logits - tf.reshape(logits_max, [-1, 1])
    sims_probs = tf.nn.softmax(sims_logits)
