def compute_hits_at_k(
    logits, labels, k):
  """Compute hits@1.

  Args:
    logits: batch_size, num_candidate
    labels: batch_size, num_candidate
    k: top k

  Returns:
    hits@k
  """
  assert logits.shape.as_list() == labels.shape.as_list()
  num_dimensions = len(logits.shape)
  _, topk = tf.nn.top_k(logits, k=k)
  hits_at_k = tf.reduce_sum(
      tf.gather(labels, topk, batch_dims=num_dimensions - 1),
      axis=1) > 0
  hits_at_k = tf.cast(hits_at_k, tf.float32)
  return hits_at_k
Compute hits@1.

  Args:
    logits: batch_size, num_candidate
    labels: batch_size, num_candidate
    k: top k

  Returns:
    hits@k
  
language/emql/util.py