add xfails for arbitrary batch sizes on some kernels

Author: pmeier

test/prototype_transforms_dispatcher_infos.py

@@ -127,6 +127,23 @@ def fill_sequence_needs_broadcast(args_kwargs):
 )
 
 
+def xfail_all_tests(*, reason, condition):
+    return [
+        TestMark(("TestDispatchers", test_name), pytest.mark.xfail(reason=reason), condition=condition)
+        for test_name in [
+            "test_scripted_smoke",
+            "test_dispatch_simple_tensor",
+            "test_dispatch_feature",
+        ]
+    ]
+
+
+xfails_degenerate_or_multi_batch_dims = xfail_all_tests(
+    reason="See https://github.com/pytorch/vision/issues/6670 for details.",
+    condition=lambda args_kwargs: len(args_kwargs.args[0].shape) > 4 or not all(args_kwargs.args[0].shape[:-3]),
+)
+
+
 DISPATCHER_INFOS = [
     DispatcherInfo(
         F.horizontal_flip,
@@ -243,6 +260,7 @@ def fill_sequence_needs_broadcast(args_kwargs):
         pil_kernel_info=PILKernelInfo(F.perspective_image_pil),
         test_marks=[
             xfail_dispatch_pil_if_fill_sequence_needs_broadcast,
+            *xfails_degenerate_or_multi_batch_dims,
         ],
     ),
     DispatcherInfo(
@@ -253,6 +271,7 @@ def fill_sequence_needs_broadcast(args_kwargs):
             features.Mask: F.elastic_mask,
         },
         pil_kernel_info=PILKernelInfo(F.elastic_image_pil),
+        test_marks=xfails_degenerate_or_multi_batch_dims,
     ),
     DispatcherInfo(
         F.center_crop,
@@ -275,6 +294,7 @@ def fill_sequence_needs_broadcast(args_kwargs):
         test_marks=[
             xfail_jit_python_scalar_arg("kernel_size"),
             xfail_jit_python_scalar_arg("sigma"),
+            *xfails_degenerate_or_multi_batch_dims,
         ],
     ),
     DispatcherInfo(
@@ -283,6 +303,7 @@ def fill_sequence_needs_broadcast(args_kwargs):
             features.Image: F.equalize_image_tensor,
         },
         pil_kernel_info=PILKernelInfo(F.equalize_image_pil, kernel_name="equalize_image_pil"),
+        test_marks=xfails_degenerate_or_multi_batch_dims,
     ),
     DispatcherInfo(
         F.invert,
@@ -318,6 +339,15 @@ def fill_sequence_needs_broadcast(args_kwargs):
             features.Image: F.adjust_sharpness_image_tensor,
         },
         pil_kernel_info=PILKernelInfo(F.adjust_sharpness_image_pil, kernel_name="adjust_sharpness_image_pil"),
+        test_marks=xfail_all_tests(
+            reason="See https://github.com/pytorch/vision/issues/6670 for details.",
+            condition=lambda args_kwargs: all(dim > 2 for dim in args_kwargs.args[0].shape[-2:])
+            and (
+                len(args_kwargs.args[0].shape) > 4
+                or not all(args_kwargs.args[0].shape[:-4])
+                or args_kwargs.args[0].shape[-4:-2] == (0, 3)
+            ),
+        ),
     ),
     DispatcherInfo(
         F.erase,

test/prototype_transforms_kernel_infos.py

@@ -143,6 +143,25 @@ def xfail_jit_list_of_ints(name, *, reason=None):
     )
 
 
+def xfail_all_tests(*, reason, condition):
+    return [
+        TestMark(("TestKernels", test_name), pytest.mark.xfail(reason=reason), condition=condition)
+        for test_name in [
+            "test_scripted_vs_eager",
+            "test_batched_vs_single",
+            "test_no_inplace",
+            "test_cuda_vs_cpu",
+            "test_dtype_and_device_consistency",
+        ]
+    ]
+
+
+xfails_image_degenerate_or_multi_batch_dims = xfail_all_tests(
+    reason="See https://github.com/pytorch/vision/issues/6670 for details.",
+    condition=lambda args_kwargs: len(args_kwargs.args[0].shape) > 4 or not all(args_kwargs.args[0].shape[:-3]),
+)
+
+
 KERNEL_INFOS = []
 
 
@@ -1093,11 +1112,7 @@ def sample_inputs_pad_video():
 
 
 def sample_inputs_perspective_image_tensor():
-    for image_loader in make_image_loaders(
-        sizes=["random"],
-        # FIXME: kernel should support arbitrary batch sizes
-        extra_dims=[(), (4,)],
-    ):
+    for image_loader in make_image_loaders(sizes=["random"]):
         for fill in [None, 128.0, 128, [12.0], [12.0 + c for c in range(image_loader.num_channels)]]:
             yield ArgsKwargs(image_loader, fill=fill, perspective_coeffs=_PERSPECTIVE_COEFFS[0])
 
@@ -1117,11 +1132,7 @@ def sample_inputs_perspective_bounding_box():
 
 
 def sample_inputs_perspective_mask():
-    for mask_loader in make_mask_loaders(
-        sizes=["random"],
-        # FIXME: kernel should support arbitrary batch sizes
-        extra_dims=[(), (4,)],
-    ):
+    for mask_loader in make_mask_loaders(sizes=["random"]):
         yield ArgsKwargs(mask_loader, perspective_coeffs=_PERSPECTIVE_COEFFS[0])
 
 
@@ -1145,6 +1156,7 @@ def sample_inputs_perspective_video():
             reference_fn=pil_reference_wrapper(F.perspective_image_pil),
             reference_inputs_fn=reference_inputs_perspective_image_tensor,
             closeness_kwargs=DEFAULT_IMAGE_CLOSENESS_KWARGS,
+            test_marks=xfails_image_degenerate_or_multi_batch_dims,
         ),
         KernelInfo(
             F.perspective_bounding_box,
@@ -1156,6 +1168,7 @@ def sample_inputs_perspective_video():
             reference_fn=pil_reference_wrapper(F.perspective_image_pil),
             reference_inputs_fn=reference_inputs_perspective_mask,
             closeness_kwargs=DEFAULT_IMAGE_CLOSENESS_KWARGS,
+            test_marks=xfails_image_degenerate_or_multi_batch_dims,
         ),
         KernelInfo(
             F.perspective_video,
@@ -1170,11 +1183,7 @@ def _get_elastic_displacement(image_size):
 
 
 def sample_inputs_elastic_image_tensor():
-    for image_loader in make_image_loaders(
-        sizes=["random"],
-        # FIXME: kernel should support arbitrary batch sizes
-        extra_dims=[(), (4,)],
-    ):
+    for image_loader in make_image_loaders(sizes=["random"]):
         displacement = _get_elastic_displacement(image_loader.image_size)
         for fill in [None, 128.0, 128, [12.0], [12.0 + c for c in range(image_loader.num_channels)]]:
             yield ArgsKwargs(image_loader, displacement=displacement, fill=fill)
@@ -1205,11 +1214,7 @@ def sample_inputs_elastic_bounding_box():
 
 
 def sample_inputs_elastic_mask():
-    for mask_loader in make_mask_loaders(
-        sizes=["random"],
-        # FIXME: kernel should support arbitrary batch sizes
-        extra_dims=[(), (4,)],
-    ):
+    for mask_loader in make_mask_loaders(sizes=["random"]):
         displacement = _get_elastic_displacement(mask_loader.shape[-2:])
         yield ArgsKwargs(mask_loader, displacement=displacement)
 
@@ -1234,6 +1239,7 @@ def sample_inputs_elastic_video():
             reference_fn=pil_reference_wrapper(F.elastic_image_pil),
             reference_inputs_fn=reference_inputs_elastic_image_tensor,
             closeness_kwargs=DEFAULT_IMAGE_CLOSENESS_KWARGS,
+            test_marks=xfails_image_degenerate_or_multi_batch_dims,
         ),
         KernelInfo(
             F.elastic_bounding_box,
@@ -1245,6 +1251,7 @@ def sample_inputs_elastic_video():
             reference_fn=pil_reference_wrapper(F.elastic_image_pil),
             reference_inputs_fn=reference_inputs_elastic_mask,
             closeness_kwargs=DEFAULT_IMAGE_CLOSENESS_KWARGS,
+            test_marks=xfails_image_degenerate_or_multi_batch_dims,
         ),
         KernelInfo(
             F.elastic_video,
@@ -1346,11 +1353,7 @@ def sample_inputs_center_crop_video():
 
 def sample_inputs_gaussian_blur_image_tensor():
     make_gaussian_blur_image_loaders = functools.partial(
-        make_image_loaders,
-        sizes=["random"],
-        color_spaces=[features.ColorSpace.RGB],
-        # FIXME: kernel should support arbitrary batch sizes
-        extra_dims=[(), (4,)],
+        make_image_loaders, sizes=["random"], color_spaces=[features.ColorSpace.RGB]
     )
 
     for image_loader, kernel_size in itertools.product(make_gaussian_blur_image_loaders(), [5, (3, 3), [3, 3]]):
@@ -1376,6 +1379,7 @@ def sample_inputs_gaussian_blur_video():
             test_marks=[
                 xfail_jit_python_scalar_arg("kernel_size"),
                 xfail_jit_python_scalar_arg("sigma"),
+                *xfails_image_degenerate_or_multi_batch_dims,
             ],
         ),
         KernelInfo(
@@ -1388,11 +1392,7 @@ def sample_inputs_gaussian_blur_video():
 
 def sample_inputs_equalize_image_tensor():
     for image_loader in make_image_loaders(
-        sizes=["random"],
-        # FIXME: kernel should support arbitrary batch sizes
-        extra_dims=[(), (4,)],
-        color_spaces=(features.ColorSpace.GRAY, features.ColorSpace.RGB),
-        dtypes=[torch.uint8],
+        sizes=["random"], color_spaces=(features.ColorSpace.GRAY, features.ColorSpace.RGB), dtypes=[torch.uint8]
     ):
         yield ArgsKwargs(image_loader)
 
@@ -1418,6 +1418,7 @@ def sample_inputs_equalize_video():
             reference_fn=pil_reference_wrapper(F.equalize_image_pil),
             reference_inputs_fn=reference_inputs_equalize_image_tensor,
             closeness_kwargs=DEFAULT_IMAGE_CLOSENESS_KWARGS,
+            test_marks=xfails_image_degenerate_or_multi_batch_dims,
         ),
         KernelInfo(
             F.equalize_video,
@@ -1594,8 +1595,6 @@ def sample_inputs_adjust_sharpness_image_tensor():
     for image_loader in make_image_loaders(
         sizes=["random", (2, 2)],
         color_spaces=(features.ColorSpace.GRAY, features.ColorSpace.RGB),
-        # FIXME: kernel should support arbitrary batch sizes
-        extra_dims=[(), (4,)],
     ):
         yield ArgsKwargs(image_loader, sharpness_factor=_ADJUST_SHARPNESS_FACTORS[0])
 
@@ -1622,6 +1621,15 @@ def sample_inputs_adjust_sharpness_video():
             reference_fn=pil_reference_wrapper(F.adjust_sharpness_image_pil),
             reference_inputs_fn=reference_inputs_adjust_sharpness_image_tensor,
             closeness_kwargs=DEFAULT_IMAGE_CLOSENESS_KWARGS,
+            test_marks=xfail_all_tests(
+                reason="See https://github.com/pytorch/vision/issues/6670 for details.",
+                condition=lambda args_kwargs: all(dim > 2 for dim in args_kwargs.args[0].shape[-2:])
+                and (
+                    len(args_kwargs.args[0].shape) > 4
+                    or not all(args_kwargs.args[0].shape[:-4])
+                    or args_kwargs.args[0].shape[-4:-2] == (0, 3)
+                ),
+            ),
         ),
         KernelInfo(
             F.adjust_sharpness_video,