def compute_logits(self, candidates, query,
                     hidden_size):
    """Compute logits between query embedding and candidate embeddings.

    Args:
      candidates: candidate embeddings
      query: query embeddings
      hidden_size: hidden size

    Returns:
      L2 logits and MIPS logits
    """

    l2_logits = - tf.reduce_sum(
        tf.square(candidates - query), axis=2, keepdims=True)
    # batch_size, num_candidate, 1
    l2_logits = l2_logits / tf.sqrt(float(hidden_size))
    # batch_size, num_candidate, 1
    l2_logits = tf.squeeze(l2_logits, axis=2)
    # batch_size, num_candidate

    mips_logits = tf.matmul(
        candidates, query, transpose_b=True) / tf.sqrt(float(hidden_size))
    # batch_size, num_candidate, 1
    mips_logits = tf.squeeze(mips_logits, axis=2)
    # batch_size, num_candidate
    return mips_logits, l2_logits
Compute logits between query embedding and candidate embeddings.

    Args:
      candidates: candidate embeddings
      query: query embeddings
      hidden_size: hidden size

    Returns:
      L2 logits and MIPS logits
    
language/emql/model.py