code: def non_max_suppression_mr(
      scores: np.ndarray,
      segments: np.ndarray,
      config: ml_collections.ConfigDict,
  ) -> Tuple[List[Array], List[Array]]:
      """Performs class-agnostic non-maximal suppression for each caption.
  
      Args:
        scores: Predicted class scores in shape (num_captions, N).
        segments: Predicted segments in shape (num_captions, N, 2).
        config: NMS configs.
  
      Returns:
        A List of scores in shape (M, ) and segments in shape (M, 2) after NMS,
        where 0 <= M <= `max_detections`.
      """
  
      out_scores = []
      out_segments = []
      num_captions = scores.shape[0]
      for i in range(num_captions):
          (
              selected_indices,
              selected_scores,
          ) = tf.image.non_max_suppression_with_scores(
              make_2d_boxes(segments[i]),
              scores[i],
              max_output_size=config.get("max_detections", 100),
              iou_threshold=config.get("iou_threshold", 0.5),
              score_threshold=config.get("score_threshold", 0.001),
              soft_nms_sigma=config.get("soft_nms_sigma", 0.3),
          )
          out_scores.append(selected_scores.numpy())
          out_segments.append(segments[i][selected_indices])
  
      return out_scores, out_segments
  
file_path: scenic/scenic/projects/unloc/postprocessing_utils.py
imports: from typing import Any, List, Optional, Tuple, Union
  import jax.numpy
  import ml_collections
  import numpy
  from scenic.projects.unloc import model_utils
  import tensorflow
label_desc: Use image node_id 5 for calling a global function with example usage: non_max_suppression_mr(scores, segments, config) and returns: out_scores,  out_segments