def compute_tf_loss(self, logits, labels):
    """Compute loss between logits and labels.

    Args:
      logits: batch_size, num_candidate
      labels: batch_size, num_candidate

    Returns:
      loss and train_op
    """
    labels = tf.cast(labels, tf.float32)
    # If labels contains all zeros, replace them with all ones
    is_label_all_zeros = tf.cast(tf.reduce_all(
        tf.equal(labels, 0.0), axis=1, keepdims=True), dtype=tf.float32)
    # batch_size, 1
    padded_labels = tf.tile(is_label_all_zeros, [1, labels.shape[1]])
    # batch_size, num_candidate
    labels += padded_labels
    # batch_size, num_candidate

    # Also zero out logits if their labels are all zeros.
    logits *= (1 - is_label_all_zeros)
    # batch_size, num_candidate

    labels_sum = tf.reduce_sum(labels, axis=1, keepdims=True)
    labels = labels / labels_sum
    loss = tf.nn.softmax_cross_entropy_with_logits(
        labels=labels, logits=logits)
    loss = tf.reduce_mean(loss)

    return loss
Compute loss between logits and labels.

    Args:
      logits: batch_size, num_candidate
      labels: batch_size, num_candidate

    Returns:
      loss and train_op
    
language/emql/model.py