def _ref_non_maximum_suppression(
        boxes, scores, iou_threshold, score_threshold, max_boxes, per_class_suppression
    ):
        """
        Reference implementation of Core ML's NMS op using TensorFlow.
        boxes of shape (n_batch, n_box, 4), [center_w, center_h, width, height]
        scores of shape (n_batch, n_box, n_score)
        output shapes [
           (n_batch, max_boxes, 4),
           (n_batch, max_boxes, n_score),
           (n_batch, max_boxes),
           (n_batch,)
        ]
        """
        n_batch, n_box, n_score = scores.shape

        iou_threshold = iou_threshold.astype(np.float32)
        score_threshold = score_threshold.astype(np.float32)

        # convert box ids to TF style
        center_w, center_h, width, height = np.split(boxes, 4, axis=-1)  # (n_batch,n_box,1)
        y1 = center_h - 0.5 * height
        y2 = center_h + 0.5 * height
        x1 = center_w - 0.5 * width
        x2 = center_w + 0.5 * width
        boxes_tf = np.concatenate((y1, x1, y2, x2), axis=-1)  # (n_batch,n_box,4)

        out1 = np.zeros((n_batch, max_boxes, 4))
        out2 = np.zeros((n_batch, max_boxes, n_score))
        out3 = -1 * np.ones((n_batch, max_boxes))
        out4 = np.zeros((n_batch,))

        for b in range(n_batch):
            box_coord_matrix = boxes_tf[b, :, :]  # (n_box,4)
            score_vector = np.max(scores[b, :, :], axis=-1)  # (n_box,)
            if not per_class_suppression:
                # this is the simple case as TF directly supports it
                ids_g = tf.image.non_max_suppression(
                    box_coord_matrix,
                    score_vector,
                    max_output_size=max_boxes,
                    iou_threshold=iou_threshold,
                    score_threshold=score_threshold,
                )
                ids = ids_g.numpy()
            else:
                # this is slightly complicated as TF does not directly support it
                class_ids = np.argmax(scores[b, :, :], axis=-1)  # (n_box,)
                sorted_score_ids = np.argsort(-score_vector)
                box_coord_matrix2 = np.take(box_coord_matrix, sorted_score_ids, axis=0)
                score_vector2 = np.take(score_vector, sorted_score_ids)
                class_ids = np.take(class_ids, sorted_score_ids)
                classes_seen = dict()
                ids_intermediate = np.array([], dtype=np.int32)
                for n in range(n_box):
                    if class_ids[n] in classes_seen:
                        continue
                    c = class_ids[n]
                    classes_seen[c] = True
                    current_class_ids = np.where(class_ids == c)[0]
                    if len(current_class_ids) > 0:
                        feed_in1 = np.take(box_coord_matrix2, current_class_ids, axis=0)
                        feed_in2 = np.take(score_vector2, current_class_ids)
                        cur_ids_g = tf.image.non_max_suppression(
                            feed_in1,
                            feed_in2,
                            max_output_size=max_boxes,
                            iou_threshold=iou_threshold,
                            score_threshold=score_threshold,
                        )
                        cur_ids = cur_ids_g.numpy()

                        from_sort_ids = np.take(current_class_ids, cur_ids)
                        ids_intermediate = np.append(ids_intermediate, from_sort_ids)
                ids_intermediate.sort()
                ids = np.take(sorted_score_ids, ids_intermediate)

            xx = len(ids)
            if xx == 0:
                ids = np.array([np.argmax(score_vector)])
                xx = 1
            if xx > max_boxes:
                ids = ids[:max_boxes]
                xx = len(ids)
            out1[b, :xx, :] = np.take(boxes[b, :, :], ids, axis=0)
            out2[b, :xx, :] = np.take(scores[b, :, :], ids, axis=0)
            out3[b, :xx] = ids
            out4[b] = xx

        return out1, out2, out3, out4

        Reference implementation of Core ML's NMS op using TensorFlow.
        boxes of shape (n_batch, n_box, 4), [center_w, center_h, width, height]
        scores of shape (n_batch, n_box, n_score)
        output shapes [
           (n_batch, max_boxes, 4),
           (n_batch, max_boxes, n_score),
           (n_batch, max_boxes),
           (n_batch,)
        ]
        
coremltools/converters/mil/mil/ops/tests/iOS14/test_tensor_operation.py