code: def compute_top_k_and_ndcg(
      logits,  # type: tf.Tensor
      duplicate_mask,  # type: tf.Tensor
      match_mlperf=False,  # type: bool
  ):
      """Compute inputs of metric calculation.
  
      Args:
        logits: A tensor containing the predicted logits for each user. The shape
          of logits is (num_users_per_batch * (1 + NUM_EVAL_NEGATIVES),) Logits
          for a user are grouped, and the first element of the group is the true
          element.
        duplicate_mask: A vector with the same shape as logits, with a value of 1
          if the item corresponding to the logit at that position has already
          appeared for that user.
        match_mlperf: Use the MLPerf reference convention for computing rank.
  
      Returns:
        is_top_k, ndcg and weights, all of which has size (num_users_in_batch,), and
        logits_by_user which has size
        (num_users_in_batch, (rconst.NUM_EVAL_NEGATIVES + 1)).
      """
      logits_by_user = tf.reshape(
          logits, (-1, rconst.NUM_EVAL_NEGATIVES + 1)
      )
      duplicate_mask_by_user = tf.cast(
          tf.reshape(
              duplicate_mask, (-1, rconst.NUM_EVAL_NEGATIVES + 1)
          ),
          logits_by_user.dtype,
      )
  
      if match_mlperf:
          # Set duplicate logits to the min value for that dtype. The MLPerf
          # reference dedupes during evaluation.
          logits_by_user *= 1 - duplicate_mask_by_user
          logits_by_user += (
              duplicate_mask_by_user * logits_by_user.dtype.min
          )
  
      # Determine the location of the first element in each row after the elements
      # are sorted.
      sort_indices = tf.argsort(
          logits_by_user, axis=1, direction="DESCENDING"
      )
  
      # Use matrix multiplication to extract the position of the true item from the
      # tensor of sorted indices. This approach is chosen because both GPUs and TPUs
      # perform matrix multiplications very quickly. This is similar to np.argwhere.
      # However this is a special case because the target will only appear in
      # sort_indices once.
      one_hot_position = tf.cast(
          tf.equal(sort_indices, rconst.NUM_EVAL_NEGATIVES), tf.int32
      )
      sparse_positions = tf.multiply(
          one_hot_position,
          tf.range(logits_by_user.shape[1])[tf.newaxis, :],
      )
      position_vector = tf.reduce_sum(sparse_positions, axis=1)
  
      in_top_k = tf.cast(
          tf.less(position_vector, rconst.TOP_K), tf.float32
      )
      ndcg = tf.math.log(2.0) / tf.math.log(
          tf.cast(position_vector, tf.float32) + 2
      )
      ndcg *= in_top_k
  
      # If a row is a padded row, all but the first element will be a duplicate.
      metric_weights = tf.not_equal(
          tf.reduce_sum(duplicate_mask_by_user, axis=1),
          rconst.NUM_EVAL_NEGATIVES,
      )
  
      return in_top_k, ndcg, metric_weights, logits_by_user
  
file_path: autodist/examples/benchmark/utils/recommendation/neumf_model.py
imports: from __future__ import absolute_import
  from __future__ import division
  from __future__ import print_function
  import sys
  from six.moves import xrange
  import tensorflow
  from utils.recommendation import constants
  from utils.recommendation import movielens
  from utils.recommendation import ncf_common
  from utils.recommendation import stat_utils
  from utils.logs import mlperf_helper
label_desc: Use image node_id 7 for calling a global function with example usage: compute_top_k_and_ndcg(logits, duplicate_mask, match_mlperf) and returns: in_top_k,  ndcg,  metric_weights,  logits_by_user