diversify parameter types for a couple of prototype kernels (#6635)

* add more size types for prototype resize sample inputs

* add skip for dispatcher

* add more sizes to resize kernel info

* add more skips

* add more diversity to gaussian_blur parameters

* diversify affine parameters and fix bounding box kernel

* fix center_crop dispatcher info

* revert kernel fixes

* add skips for scalar shears in affine_bounding_box

Author: pmeier

test/prototype_transforms_dispatcher_infos.py

@@ -1,5 +1,6 @@
 import dataclasses
-from typing import Callable, Dict, Sequence, Type
+from collections import defaultdict
+from typing import Callable, Dict, List, Sequence, Type
 
 import pytest
 import torchvision.prototype.transforms.functional as F
@@ -11,15 +12,30 @@
 KERNEL_SAMPLE_INPUTS_FN_MAP = {info.kernel: info.sample_inputs_fn for info in KERNEL_INFOS}
 
 
+def skip_python_scalar_arg_jit(name, *, reason="Python scalar int or float is not supported when scripting"):
+    return Skip(
+        "test_scripted_smoke",
+        condition=lambda args_kwargs, device: isinstance(args_kwargs.kwargs[name], (int, float)),
+        reason=reason,
+    )
+
+
+def skip_integer_size_jit(name="size"):
+    return skip_python_scalar_arg_jit(name, reason="Integer size is not supported when scripting.")
+
+
 @dataclasses.dataclass
 class DispatcherInfo:
     dispatcher: Callable
     kernels: Dict[Type, Callable]
     skips: Sequence[Skip] = dataclasses.field(default_factory=list)
-    _skips_map: Dict[str, Skip] = dataclasses.field(default=None, init=False)
+    _skips_map: Dict[str, List[Skip]] = dataclasses.field(default=None, init=False)
 
     def __post_init__(self):
-        self._skips_map = {skip.test_name: skip for skip in self.skips}
+        skips_map = defaultdict(list)
+        for skip in self.skips:
+            skips_map[skip.test_name].append(skip)
+        self._skips_map = dict(skips_map)
 
     def sample_inputs(self, *types):
         for type in types or self.kernels.keys():
@@ -29,9 +45,13 @@ def sample_inputs(self, *types):
             yield from KERNEL_SAMPLE_INPUTS_FN_MAP[self.kernels[type]]()
 
     def maybe_skip(self, *, test_name, args_kwargs, device):
-        skip = self._skips_map.get(test_name)
-        if skip and skip.condition(args_kwargs, device):
-            pytest.skip(skip.reason)
+        skips = self._skips_map.get(test_name)
+        if not skips:
+            return
+
+        for skip in skips:
+            if skip.condition(args_kwargs, device):
+                pytest.skip(skip.reason)
 
 
 DISPATCHER_INFOS = [
@@ -50,6 +70,9 @@ def maybe_skip(self, *, test_name, args_kwargs, device):
             features.BoundingBox: F.resize_bounding_box,
             features.Mask: F.resize_mask,
         },
+        skips=[
+            skip_integer_size_jit(),
+        ],
     ),
     DispatcherInfo(
         F.affine,
@@ -58,6 +81,7 @@ def maybe_skip(self, *, test_name, args_kwargs, device):
             features.BoundingBox: F.affine_bounding_box,
             features.Mask: F.affine_mask,
         },
+        skips=[skip_python_scalar_arg_jit("shear", reason="Scalar shear is not supported by JIT")],
     ),
     DispatcherInfo(
         F.vertical_flip,
@@ -122,12 +146,19 @@ def maybe_skip(self, *, test_name, args_kwargs, device):
             features.BoundingBox: F.center_crop_bounding_box,
             features.Mask: F.center_crop_mask,
         },
+        skips=[
+            skip_integer_size_jit("output_size"),
+        ],
     ),
     DispatcherInfo(
         F.gaussian_blur,
         kernels={
             features.Image: F.gaussian_blur_image_tensor,
         },
+        skips=[
+            skip_python_scalar_arg_jit("kernel_size"),
+            skip_python_scalar_arg_jit("sigma"),
+        ],
     ),
     DispatcherInfo(
         F.equalize,
@@ -207,11 +238,7 @@ def maybe_skip(self, *, test_name, args_kwargs, device):
             features.Image: F.five_crop_image_tensor,
         },
         skips=[
-            Skip(
-                "test_scripted_smoke",
-                condition=lambda args_kwargs, device: isinstance(args_kwargs.kwargs["size"], int),
-                reason="Integer size is not supported when scripting five_crop_image_tensor.",
-            ),
+            skip_integer_size_jit(),
         ],
     ),
     DispatcherInfo(
@@ -220,11 +247,7 @@ def maybe_skip(self, *, test_name, args_kwargs, device):
             features.Image: F.ten_crop_image_tensor,
         },
         skips=[
-            Skip(
-                "test_scripted_smoke",
-                condition=lambda args_kwargs, device: isinstance(args_kwargs.kwargs["size"], int),
-                reason="Integer size is not supported when scripting ten_crop_image_tensor.",
-            ),
+            skip_integer_size_jit(),
         ],
     ),
     DispatcherInfo(