[fbsync] Dispatcher -> Functional (#7829)

Reviewed By: matteobettini

Differential Revision: D48642298

fbshipit-source-id: fc7c521bde09fea9450b5ac2b11b5b055432db65

Author: NicolasHug

gallery/plot_custom_datapoints.py

@@ -49,7 +49,7 @@ class MyDatapoint(datapoints.Datapoint):
 from torchvision.transforms.v2 import functional as F
 
 
-@F.register_kernel(dispatcher="hflip", datapoint_cls=MyDatapoint)
+@F.register_kernel(functional="hflip", datapoint_cls=MyDatapoint)
 def hflip_my_datapoint(my_dp, *args, **kwargs):
     print("Flipping!")
     out = my_dp.flip(-1)
@@ -64,9 +64,9 @@ def hflip_my_datapoint(my_dp, *args, **kwargs):
 # .. note::
 #
 #     In our call to ``register_kernel`` above we used a string
-#     ``dispatcher="hflip"`` to refer to the functional we want to hook into. We
+#     ``functional="hflip"`` to refer to the functional we want to hook into. We
 #     could also have used the  functional *itself*, i.e.
-#     ``@register_kernel(dispatcher=F.hflip, ...)``.
+#     ``@register_kernel(functional=F.hflip, ...)``.
 #
 #     The functionals that you can be hooked into are the ones in
 #     ``torchvision.transforms.v2.functional`` and they are documented in

test/test_transforms_v2_refactored.py

@@ -91,13 +91,13 @@ def __init__(
 
         self._log_ratio = torch.log(torch.tensor(self.ratio))
 
-    def _call_kernel(self, dispatcher: Callable, inpt: Any, *args: Any, **kwargs: Any) -> Any:
+    def _call_kernel(self, functional: Callable, inpt: Any, *args: Any, **kwargs: Any) -> Any:
         if isinstance(inpt, (datapoints.BoundingBoxes, datapoints.Mask)):
             warnings.warn(
                 f"{type(self).__name__}() is currently passing through inputs of type "
                 f"datapoints.{type(inpt).__name__}. This will likely change in the future."
             )
-        return super()._call_kernel(dispatcher, inpt, *args, **kwargs)
+        return super()._call_kernel(functional, inpt, *args, **kwargs)
 
     def _get_params(self, flat_inputs: List[Any]) -> Dict[str, Any]:
         img_c, img_h, img_w = query_chw(flat_inputs)

torchvision/transforms/v2/_augment.py

@@ -358,13 +358,13 @@ def __init__(self, size: Union[int, Sequence[int]]) -> None:
         super().__init__()
         self.size = _setup_size(size, error_msg="Please provide only two dimensions (h, w) for size.")
 
-    def _call_kernel(self, dispatcher: Callable, inpt: Any, *args: Any, **kwargs: Any) -> Any:
+    def _call_kernel(self, functional: Callable, inpt: Any, *args: Any, **kwargs: Any) -> Any:
         if isinstance(inpt, (datapoints.BoundingBoxes, datapoints.Mask)):
             warnings.warn(
                 f"{type(self).__name__}() is currently passing through inputs of type "
                 f"datapoints.{type(inpt).__name__}. This will likely change in the future."
             )
-        return super()._call_kernel(dispatcher, inpt, *args, **kwargs)
+        return super()._call_kernel(functional, inpt, *args, **kwargs)
 
     def _transform(self, inpt: Any, params: Dict[str, Any]) -> Any:
         return self._call_kernel(F.five_crop, inpt, self.size)
@@ -405,13 +405,13 @@ def __init__(self, size: Union[int, Sequence[int]], vertical_flip: bool = False)
         self.size = _setup_size(size, error_msg="Please provide only two dimensions (h, w) for size.")
         self.vertical_flip = vertical_flip
 
-    def _call_kernel(self, dispatcher: Callable, inpt: Any, *args: Any, **kwargs: Any) -> Any:
+    def _call_kernel(self, functional: Callable, inpt: Any, *args: Any, **kwargs: Any) -> Any:
         if isinstance(inpt, (datapoints.BoundingBoxes, datapoints.Mask)):
             warnings.warn(
                 f"{type(self).__name__}() is currently passing through inputs of type "
                 f"datapoints.{type(inpt).__name__}. This will likely change in the future."
             )
-        return super()._call_kernel(dispatcher, inpt, *args, **kwargs)
+        return super()._call_kernel(functional, inpt, *args, **kwargs)
 
     def _check_inputs(self, flat_inputs: List[Any]) -> None:
         if has_any(flat_inputs, datapoints.BoundingBoxes, datapoints.Mask):

torchvision/transforms/v2/_geometry.py

@@ -30,8 +30,8 @@ def _check_inputs(self, flat_inputs: List[Any]) -> None:
     def _get_params(self, flat_inputs: List[Any]) -> Dict[str, Any]:
         return dict()
 
-    def _call_kernel(self, dispatcher: Callable, inpt: Any, *args: Any, **kwargs: Any) -> Any:
-        kernel = _get_kernel(dispatcher, type(inpt), allow_passthrough=True)
+    def _call_kernel(self, functional: Callable, inpt: Any, *args: Any, **kwargs: Any) -> Any:
+        kernel = _get_kernel(functional, type(inpt), allow_passthrough=True)
         return kernel(inpt, *args, **kwargs)
 
     def _transform(self, inpt: Any, params: Dict[str, Any]) -> Any:

torchvision/transforms/v2/_transform.py

@@ -203,7 +203,7 @@ def convert_format_bounding_boxes(
     new_format: Optional[BoundingBoxFormat] = None,
     inplace: bool = False,
 ) -> torch.Tensor:
-    # This being a kernel / dispatcher hybrid, we need an option to pass `old_format` explicitly for simple tensor
+    # This being a kernel / functional hybrid, we need an option to pass `old_format` explicitly for simple tensor
     # inputs as well as extract it from `datapoints.BoundingBoxes` inputs. However, putting a default value on
     # `old_format` means we also need to put one on `new_format` to have syntactically correct Python. Here we mimic the
     # default error that would be thrown if `new_format` had no default value.

torchvision/transforms/v2/functional/_meta.py

@@ -12,7 +12,7 @@ def is_simple_tensor(inpt: Any) -> bool:
     return isinstance(inpt, torch.Tensor) and not isinstance(inpt, datapoints.Datapoint)
 
 
-# {dispatcher: {input_type: type_specific_kernel}}
+# {functional: {input_type: type_specific_kernel}}
 _KERNEL_REGISTRY: Dict[Callable, Dict[Type, Callable]] = {}
 
 
@@ -27,10 +27,10 @@ def wrapper(inpt, *args, **kwargs):
     return wrapper
 
 
-def _register_kernel_internal(dispatcher, input_type, *, datapoint_wrapper=True):
-    registry = _KERNEL_REGISTRY.setdefault(dispatcher, {})
+def _register_kernel_internal(functional, input_type, *, datapoint_wrapper=True):
+    registry = _KERNEL_REGISTRY.setdefault(functional, {})
     if input_type in registry:
-        raise ValueError(f"Dispatcher {dispatcher} already has a kernel registered for type {input_type}.")
+        raise ValueError(f"Functional {functional} already has a kernel registered for type {input_type}.")
 
     def decorator(kernel):
         registry[input_type] = (
@@ -43,14 +43,14 @@ def decorator(kernel):
     return decorator
 
 
-def _name_to_dispatcher(name):
+def _name_to_functional(name):
     import torchvision.transforms.v2.functional  # noqa
 
     try:
         return getattr(torchvision.transforms.v2.functional, name)
     except AttributeError:
         raise ValueError(
-            f"Could not find dispatcher with name '{name}' in torchvision.transforms.v2.functional."
+            f"Could not find functional with name '{name}' in torchvision.transforms.v2.functional."
         ) from None
 
 
@@ -59,21 +59,21 @@ def _name_to_dispatcher(name):
 }
 
 
-def register_kernel(dispatcher, datapoint_cls):
-    """Decorate a kernel to register it for a dispatcher and a (custom) datapoint type.
+def register_kernel(functional, datapoint_cls):
+    """Decorate a kernel to register it for a functional and a (custom) datapoint type.
 
     See :ref:`sphx_glr_auto_examples_plot_custom_datapoints.py` for usage
     details.
     """
-    if isinstance(dispatcher, str):
-        dispatcher = _name_to_dispatcher(name=dispatcher)
+    if isinstance(functional, str):
+        functional = _name_to_functional(name=functional)
     elif not (
-        callable(dispatcher)
-        and getattr(dispatcher, "__module__", "").startswith("torchvision.transforms.v2.functional")
+        callable(functional)
+        and getattr(functional, "__module__", "").startswith("torchvision.transforms.v2.functional")
     ):
         raise ValueError(
-            f"Kernels can only be registered on dispatchers from the torchvision.transforms.v2.functional namespace, "
-            f"but got {dispatcher}."
+            f"Kernels can only be registered on functionals from the torchvision.transforms.v2.functional namespace, "
+            f"but got {functional}."
         )
 
     if not (isinstance(datapoint_cls, type) and issubclass(datapoint_cls, datapoints.Datapoint)):
@@ -85,13 +85,13 @@ def register_kernel(dispatcher, datapoint_cls):
     if datapoint_cls in _BUILTIN_DATAPOINT_TYPES:
         raise ValueError(f"Kernels cannot be registered for the builtin datapoint classes, but got {datapoint_cls}")
 
-    return _register_kernel_internal(dispatcher, datapoint_cls, datapoint_wrapper=False)
+    return _register_kernel_internal(functional, datapoint_cls, datapoint_wrapper=False)
 
 
-def _get_kernel(dispatcher, input_type, *, allow_passthrough=False):
-    registry = _KERNEL_REGISTRY.get(dispatcher)
+def _get_kernel(functional, input_type, *, allow_passthrough=False):
+    registry = _KERNEL_REGISTRY.get(functional)
     if not registry:
-        raise ValueError(f"No kernel registered for dispatcher {dispatcher.__name__}.")
+        raise ValueError(f"No kernel registered for functional {functional.__name__}.")
 
     # In case we have an exact type match, we take a shortcut.
     if input_type in registry:
@@ -113,17 +113,17 @@ def _get_kernel(dispatcher, input_type, *, allow_passthrough=False):
         return lambda inpt, *args, **kwargs: inpt
 
     raise TypeError(
-        f"Dispatcher F.{dispatcher.__name__} supports inputs of type {registry.keys()}, "
+        f"Functional F.{functional.__name__} supports inputs of type {registry.keys()}, "
         f"but got {input_type} instead."
     )
 
 
 # This basically replicates _register_kernel_internal, but with a specialized wrapper for five_crop / ten_crop
-# We could get rid of this by letting _register_kernel_internal take arbitrary dispatchers rather than wrap_kernel: bool
-def _register_five_ten_crop_kernel_internal(dispatcher, input_type):
-    registry = _KERNEL_REGISTRY.setdefault(dispatcher, {})
+# We could get rid of this by letting _register_kernel_internal take arbitrary functionals rather than wrap_kernel: bool
+def _register_five_ten_crop_kernel_internal(functional, input_type):
+    registry = _KERNEL_REGISTRY.setdefault(functional, {})
     if input_type in registry:
-        raise TypeError(f"Dispatcher '{dispatcher}' already has a kernel registered for type '{input_type}'.")
+        raise TypeError(f"Functional '{functional}' already has a kernel registered for type '{input_type}'.")
 
     def wrap(kernel):
         @functools.wraps(kernel)