code: def test_box_nms_op():
      def test_box_nms_forward(
          data,
          expected,
          thresh=0.5,
          valid=0,
          topk=-1,
          coord=2,
          score=1,
          cid=0,
          bid=-1,
          force=False,
          in_format="corner",
          out_format="corner",
      ):
          for dtype in ["float16", "float32", "float64"]:
              data = mx.nd.array(data, dtype=dtype)
              out = mx.contrib.nd.box_nms(
                  data,
                  overlap_thresh=thresh,
                  valid_thresh=valid,
                  topk=topk,
                  coord_start=coord,
                  score_index=score,
                  id_index=cid,
                  background_id=bid,
                  force_suppress=force,
                  in_format=in_format,
                  out_format=out_format,
              )
              assert_almost_equal(
                  out.asnumpy(),
                  expected.astype(dtype),
                  rtol=1e-3,
                  atol=1e-3,
              )
  
      def test_box_nms_backward(
          data,
          grad,
          expected,
          thresh=0.5,
          valid=0,
          topk=-1,
          coord=2,
          score=1,
          cid=0,
          bid=-1,
          force=False,
          in_format="corner",
          out_format="corner",
      ):
          in_var = mx.sym.Variable("data")
          arr_data = mx.nd.array(data)
          arr_grad = mx.nd.empty(arr_data.shape)
          op = mx.contrib.sym.box_nms(
              in_var,
              overlap_thresh=thresh,
              valid_thresh=valid,
              topk=topk,
              coord_start=coord,
              score_index=score,
              id_index=cid,
              background_id=bid,
              force_suppress=force,
              in_format=in_format,
              out_format=out_format,
          )
          exe = op._bind(
              ctx=default_device(),
              args=[arr_data],
              args_grad=[arr_grad],
          )
          exe.forward(is_train=True)
          exe.backward(mx.nd.array(grad))
          assert_almost_equal(arr_grad.asnumpy(), expected)
  
      def corner_to_center(data):
          out = np.reshape(data, (-1, 6)).copy()
          out[:, 2] = (data[:, 2] + data[:, 4]) / 2.0
          out[:, 3] = (data[:, 3] + data[:, 5]) / 2.0
          out[:, 4] = data[:, 4] - data[:, 2]
          out[:, 5] = data[:, 5] - data[:, 3]
          invalid = np.where(data[:, 0] < 0)[0]
          out[invalid, :] = -1
          return out
  
      def center_to_corner(data):
          data = np.reshape(data, (-1, 6)).copy()
          out[:, 2] = data[:, 2] - data[:, 4] / 2.0
          out[:, 3] = data[:, 3] - data[:, 5] / 2.0
          out[:, 4] = data[:, 2] + data[:, 4] / 2.0
          out[:, 5] = data[:, 3] + data[:, 5] / 2.0
          invalid = np.where(data[:, 0] < 0)[0]
          out[invalid, :] = -1
          return out
  
      def swap_position(
          data, expected, coord=2, score=1, cid=0, new_col=0
      ):
          data = np.reshape(data, (-1, 6))
          expected = np.reshape(expected, (-1, 6))
          new_coord = random.randint(0, 6 + new_col - 4)
          others = list(range(new_coord)) + list(
              range(new_coord + 4, 6 + new_col)
          )
          random.shuffle(others)
          new_score = others[0]
          new_cid = others[1]
          new_data = np.full(
              (data.shape[0], data.shape[1] + new_col), -1.0
          )
          new_expected = np.full(
              (expected.shape[0], expected.shape[1] + new_col), -1.0
          )
          new_data[:, new_coord : new_coord + 4] = data[
              :, coord : coord + 4
          ]
          new_data[:, new_score] = data[:, score]
          new_data[:, new_cid] = data[:, cid]
          new_expected[:, new_coord : new_coord + 4] = expected[
              :, coord : coord + 4
          ]
          new_expected[:, new_score] = expected[:, score]
          new_expected[:, new_cid] = expected[:, cid]
          return new_data, new_expected, new_coord, new_score, new_cid
  
      # manually set up test cases
      boxes = [
          [0, 0.5, 0.1, 0.1, 0.2, 0.2],
          [1, 0.4, 0.1, 0.1, 0.2, 0.2],
          [0, 0.3, 0.1, 0.1, 0.14, 0.14],
          [2, 0.6, 0.5, 0.5, 0.7, 0.8],
      ]
  
      # case1
      force = True
      thresh = 0.5
      expected = [
          [2, 0.6, 0.5, 0.5, 0.7, 0.8],
          [0, 0.5, 0.1, 0.1, 0.2, 0.2],
          [0, 0.3, 0.1, 0.1, 0.14, 0.14],
          [-1, -1, -1, -1, -1, -1],
      ]
      grad = np.random.rand(4, 6)
      expected_in_grad = grad[(1, 3, 2, 0), :]
      expected_in_grad[1, :] = 0
      test_box_nms_forward(
          np.array(boxes),
          np.array(expected),
          force=force,
          thresh=thresh,
      )
      test_box_nms_backward(
          np.array(boxes),
          grad,
          expected_in_grad,
          force=force,
          thresh=thresh,
      )
  
      # case2: multi batch
      boxes2 = [boxes] * 3
      expected2 = [expected] * 3
      grad2 = np.array([grad.tolist()] * 3)
      expected_in_grad2 = np.array([expected_in_grad.tolist()] * 3)
      test_box_nms_forward(
          np.array(boxes2),
          np.array(expected2),
          force=force,
          thresh=thresh,
      )
      test_box_nms_backward(
          np.array(boxes2),
          grad2,
          expected_in_grad2,
          force=force,
          thresh=thresh,
      )
      # another new dim
      boxes2 = [boxes2] * 2
      expected2 = [expected2] * 2
      grad2 = np.array([grad2.tolist()] * 2)
      expected_in_grad2 = np.array([expected_in_grad2.tolist()] * 2)
      test_box_nms_forward(
          np.array(boxes2),
          np.array(expected2),
          force=force,
          thresh=thresh,
      )
      test_box_nms_backward(
          np.array(boxes2),
          grad2,
          expected_in_grad2,
          force=force,
          thresh=thresh,
      )
  
      # case3: thresh
      thresh = 0.1
      boxes3 = boxes
      expected3 = [
          [2, 0.6, 0.5, 0.5, 0.7, 0.8],
          [0, 0.5, 0.1, 0.1, 0.2, 0.2],
          [-1, -1, -1, -1, -1, -1],
          [-1, -1, -1, -1, -1, -1],
      ]
      grad3 = np.random.rand(4, 6)
      expected_in_grad3 = grad3[(1, 3, 2, 0), :]
      expected_in_grad3[(1, 2), :] = 0
      test_box_nms_forward(
          np.array(boxes3),
          np.array(expected3),
          force=force,
          thresh=thresh,
      )
      test_box_nms_backward(
          np.array(boxes3),
          grad3,
          expected_in_grad3,
          force=force,
          thresh=thresh,
      )
  
      # case4: non-force
      boxes4 = boxes
      force = False
      expected4 = [
          [2, 0.6, 0.5, 0.5, 0.7, 0.8],
          [0, 0.5, 0.1, 0.1, 0.2, 0.2],
          [1, 0.4, 0.1, 0.1, 0.2, 0.2],
          [-1, -1, -1, -1, -1, -1],
      ]
      grad4 = np.random.rand(4, 6)
      expected_in_grad4 = grad4[(1, 2, 3, 0), :]
      expected_in_grad4[2, :] = 0
      test_box_nms_forward(
          np.array(boxes4),
          np.array(expected4),
          force=force,
          thresh=thresh,
      )
      test_box_nms_backward(
          np.array(boxes4),
          grad4,
          expected_in_grad4,
          force=force,
          thresh=thresh,
      )
  
      # case5: different coding
      boxes5 = corner_to_center(np.array(boxes4))
      test_box_nms_forward(
          np.array(boxes5),
          np.array(expected4),
          force=force,
          thresh=thresh,
          in_format="center",
      )
      expected5 = corner_to_center(np.array(expected4))
      test_box_nms_forward(
          np.array(boxes4),
          np.array(expected5),
          force=force,
          thresh=thresh,
          out_format="center",
      )
      test_box_nms_forward(
          np.array(boxes5),
          np.array(expected5),
          force=force,
          thresh=thresh,
          in_format="center",
          out_format="center",
      )
  
      # case6: different position
      boxes6, expected6, new_coord, new_score, new_id = swap_position(
          np.array(boxes4), np.array(expected4), new_col=2
      )
      test_box_nms_forward(
          np.array(boxes6),
          np.array(expected6),
          force=force,
          thresh=thresh,
          coord=new_coord,
          score=new_score,
          cid=new_id,
      )
  
      # case7: no id, should be same with force=True
      force = False
      thresh = 0.5
      test_box_nms_forward(
          np.array(boxes),
          np.array(expected),
          force=force,
          thresh=thresh,
          cid=-1,
      )
  
      # case8: multi-batch thresh + topk
      boxes8 = [
          [
              [1, 1, 0, 0, 10, 10],
              [1, 0.4, 0, 0, 10, 10],
              [1, 0.3, 0, 0, 10, 10],
          ],
          [
              [2, 1, 0, 0, 10, 10],
              [2, 0.4, 0, 0, 10, 10],
              [2, 0.3, 0, 0, 10, 10],
          ],
          [
              [3, 1, 0, 0, 10, 10],
              [3, 0.4, 0, 0, 10, 10],
              [3, 0.3, 0, 0, 10, 10],
          ],
      ]
      expected8 = [
          [
              [1, 1, 0, 0, 10, 10],
              [-1, -1, -1, -1, -1, -1],
              [-1, -1, -1, -1, -1, -1],
          ],
          [
              [2, 1, 0, 0, 10, 10],
              [-1, -1, -1, -1, -1, -1],
              [-1, -1, -1, -1, -1, -1],
          ],
          [
              [3, 1, 0, 0, 10, 10],
              [-1, -1, -1, -1, -1, -1],
              [-1, -1, -1, -1, -1, -1],
          ],
      ]
      grad8 = np.random.rand(3, 3, 6)
      expected_in_grad8 = np.zeros((3, 3, 6))
      expected_in_grad8[(0, 1, 2), (0, 0, 0), :] = grad8[
          (0, 1, 2), (0, 0, 0), :
      ]
      force = False
      thresh = 0.5
      valid = 0.5
      topk = 2
      test_box_nms_forward(
          np.array(boxes8),
          np.array(expected8),
          force=force,
          thresh=thresh,
          valid=valid,
          topk=topk,
      )
      test_box_nms_backward(
          np.array(boxes8),
          grad8,
          expected_in_grad8,
          force=force,
          thresh=thresh,
          valid=valid,
          topk=topk,
      )
  
      # case9: background id filter out
      # default background id -1
      boxes9 = [
          [0, 0.5, 0.1, 0.1, 0.2, 0.2],
          [0, 0.4, 0.1, 0.1, 0.2, 0.2],
          [1, 0.3, 0.1, 0.1, 0.14, 0.14],
          [-1, 0.6, 0.5, 0.5, 0.7, 0.8],
      ]
      expected9 = [
          [0, 0.5, 0.1, 0.1, 0.2, 0.2],
          [1, 0.3, 0.1, 0.1, 0.14, 0.14],
          [-1, -1, -1, -1, -1, -1],
          [-1, -1, -1, -1, -1, -1],
      ]
      force = True
      thresh = 0.5
      grad9 = np.random.rand(4, 6)
      expected_in_grad9 = grad9[(0, 2, 1, 3), :]
      expected_in_grad9[(1, 3), :] = 0
      test_box_nms_forward(
          np.array(boxes9),
          np.array(expected9),
          force=force,
          thresh=thresh,
      )
      test_box_nms_backward(
          np.array(boxes9),
          grad9,
          expected_in_grad9,
          force=force,
          thresh=thresh,
      )
      # set background id
      background_id = 0
      expected9 = [
          [-1, 0.6, 0.5, 0.5, 0.7, 0.8],
          [1, 0.3, 0.1, 0.1, 0.14, 0.14],
          [-1, -1, -1, -1, -1, -1],
          [-1, -1, -1, -1, -1, -1],
      ]
      grad9 = np.random.rand(4, 6)
      expected_in_grad9 = grad9[(2, 3, 1, 0), :]
      expected_in_grad9[(0, 1), :] = 0
      test_box_nms_forward(
          np.array(boxes9),
          np.array(expected9),
          force=force,
          thresh=thresh,
          bid=background_id,
      )
      test_box_nms_backward(
          np.array(boxes9),
          grad9,
          expected_in_grad9,
          force=force,
          thresh=thresh,
          bid=background_id,
      )
  
file_path: mxnet/tests/python/unittest/test_contrib_operator.py
imports: from __future__ import print_function
  import numpy
  import mxnet
  import random
  import itertools
  from numpy.testing import assert_allclose, assert_array_equal
  from mxnet.test_utils import *
  from common import assert_raises_cudnn_not_satisfied, xfail_when_nonstandard_decimal_separator
  import unittest
label_desc: Use image node_id 1 for calling a global function with example usage: test_box_nms_op() and returns: out, out, new_data,  new_expected,  new_coord,  new_score,  new_cid