code: def _index_predictions(predictions, labels):
      """
      Indexes predictions, a [batch_size, num_classes]-shaped tensor,
      by labels, a [batch_size]-shaped tensor that indicates which
      class each sample should be indexed by.
  
      Args:
          predictions: A [batch_size, num_classes]-shaped tensor. The input to a model.
          labels: A [batch_size, num_classes]-shaped tensor.
                  The tensor used to index predictions, in one-hot encoding form.
      Returns:
          A tensor of shape [batch_size] representing the predictions indexed by the labels.
      """
      current_batch_size = tf.shape(predictions)[0]
      sample_indices = tf.range(current_batch_size)
      sparse_labels = tf.argmax(labels, axis=-1)
      indices_tensor = tf.stack(
          [sample_indices, tf.cast(sparse_labels, tf.int32)], axis=1
      )
      predictions_indexed = tf.gather_nd(predictions, indices_tensor)
      return predictions_indexed
  
file_path: attributionpriors/attributionpriors/eager_ops.py
imports: import tensorflow
  import numpy
label_desc: Use image node_id 1 for calling a global function with example usage: _index_predictions(predictions, labels) and returns: predictions_indexed