import tensorflow as tf

    self.assertAllEqual(nms_masks2.numpy(), nms_masks_expected2.numpy())
    self.assertAllClose(nms_scores2.numpy(), nms_scores_expected2.numpy())
    self.assertAllEqual(nms_classes2.numpy(), nms_classes_expected2.numpy())

  def test_instance_non_maximum_suppression_2d_scores(self):
    mask0 = tf.constant([[1, 0],
                         [0, 1]], dtype=tf.float32)
    mask1 = tf.constant([[1, 1],
                         [0, 1]], dtype=tf.float32)
    mask2 = tf.constant([[1, 0],
                         [1, 1]], dtype=tf.float32)
    mask3 = tf.constant([[1, 1],
                         [1, 1]], dtype=tf.float32)
    mask4 = tf.constant([[0, 0],
