[TF:XLA] Allow fallback on tests of dynamic shapes.

PiperOrigin-RevId: 269514978
diff --git a/tensorflow/compiler/tf2xla/kernels/image_ops.cc b/tensorflow/compiler/tf2xla/kernels/image_ops.cc
index dcd523e..e7bf343 100644
--- a/tensorflow/compiler/tf2xla/kernels/image_ops.cc
+++ b/tensorflow/compiler/tf2xla/kernels/image_ops.cc
@@ -423,7 +423,7 @@
         errors::InvalidArgument("scores size must equal number of boxes",
                                 scores_shape.DebugString()));
     OP_REQUIRES(context, pad_to_max_output_size_,
-                errors::InvalidArgument(
+                errors::Unimplemented(
                     "XLA compilation requires pad_to_max_output_size == True"));
     OP_REQUIRES(context, num_boxes <= kint32max,
                 errors::InvalidArgument("XLA compilation requires number of "
diff --git a/tensorflow/python/ops/image_ops_test.py b/tensorflow/python/ops/image_ops_test.py
index 3f2513b..26cbbd9 100644
--- a/tensorflow/python/ops/image_ops_test.py
+++ b/tensorflow/python/ops/image_ops_test.py
@@ -4451,6 +4451,8 @@
       image_ops.non_max_suppression(boxes, scores, 3, [[0.5]])
 
   @test_util.run_deprecated_v1
+  @test_util.xla_allow_fallback(
+      "non_max_suppression with dynamic output shape unsupported.")
   def testDataTypes(self):
     # Test case for GitHub issue 20199.
     boxes_np = [[0, 0, 1, 1], [0, 0.1, 1, 1.1], [0, -0.1, 1, 0.9],
@@ -4514,7 +4516,7 @@
 class NonMaxSuppressionWithScoresTest(test_util.TensorFlowTestCase):
 
   @test_util.run_deprecated_v1
-  @test_util.disable_xla(
+  @test_util.xla_allow_fallback(
       "non_max_suppression with dynamic output shape unsupported.")
   def testSelectFromThreeClustersWithSoftNMS(self):
     boxes_np = [[0, 0, 1, 1], [0, 0.1, 1, 1.1], [0, -0.1, 1, 0.9],
@@ -4549,6 +4551,8 @@
 class NonMaxSuppressionPaddedTest(test_util.TensorFlowTestCase):
 
   @test_util.run_deprecated_v1
+  @test_util.xla_allow_fallback(
+      "non_max_suppression with dynamic output shape unsupported.")
   def testSelectFromThreeClusters(self):
     boxes_np = [[0, 0, 1, 1], [0, 0.1, 1, 1.1], [0, -0.1, 1, 0.9],
                 [0, 10, 1, 11], [0, 10.1, 1, 11.1], [0, 100, 1, 101]]
@@ -4582,6 +4586,8 @@
       self.assertEqual(num_valid.eval(), 3)
 
   @test_util.run_deprecated_v1
+  @test_util.xla_allow_fallback(
+      "non_max_suppression with dynamic output shape unsupported.")
   def testSelectFromContinuousOverLap(self):
     boxes_np = [[0, 0, 1, 1], [0, 0.2, 1, 1.2], [0, 0.4, 1, 1.4],
                 [0, 0.6, 1, 1.6], [0, 0.8, 1, 1.8], [0, 2, 1, 2]]