def compute_average_precision_at_k(
    logits, labels, k):
  """Compute average precision at k.

  Args:
    logits: batch_size, nun_candidate
    labels: batch_size, num_candidate
    k: scalar

  Returns:
    average_precision_at_k
  """

  _, topk = tf.nn.top_k(logits, k)
  # batch_size, k
  true_positives = tf.gather(labels, topk, batch_dims=1)
  # batch_size, k
  # e.g. [[0, 1, 1, 0], [1, 0, 0, 1]]
  upper_triangle_matrix = tf.constant(
      np.triu(np.ones([k, k])), dtype=tf.float32)
  # k, k
  # e.g. [[1,1,1,1], [0,1,1,1], [0,0,1,1], [0,0,0,1]]
  upper_triangle_matrix /= tf.reduce_sum(
      upper_triangle_matrix, axis=0, keepdims=True)
  # k, k
  # e.g. [[1,1/2,1/3,1/4], [0,1/2,1/3,1/4], [0,0,1/3,1/4], [0,0,0,1/4]]
  recall_at_k = tf.matmul(true_positives, upper_triangle_matrix)
  # batch_size, k
  # e.g. [[0, 1/2, 2/3, 2/4], [1, 1/2, 1/3, 2/4]]
  positive_recall_at_k = true_positives * recall_at_k
  # batch_size, k
  # e.g. [[0, 1/2, 2/3, 0], [1, 0, 0, 2/4]]
  num_true_positive = tf.reduce_sum(true_positives, axis=1)
  # batch_size
  # e.g. [2, 2]
  num_true_positive_replace_0_with_1 = tf.where(
      num_true_positive > 0, num_true_positive,
      tf.ones(tf.shape(num_true_positive), dtype=tf.float32))
  average_precision_at_k = \
      tf.reduce_sum(positive_recall_at_k, axis=1) \
      / num_true_positive_replace_0_with_1
  # batch_size
  # e.g. [(1/2 + 2/3) / 2, (1 + 2/4) / 2]
  return average_precision_at_k
Compute average precision at k.

  Args:
    logits: batch_size, nun_candidate
    labels: batch_size, num_candidate
    k: scalar

  Returns:
    average_precision_at_k
  
language/emql/util.py