From bf6a8dc2156b9761e7bcdd0df605cc1d875f8435 Mon Sep 17 00:00:00 2001
From: Nicolas Hug <nh.nicolas.hug@gmail.com>
Date: Thu, 10 Aug 2023 11:04:59 +0100
Subject: [PATCH] Simplify _NO_WRAPPING_EXCEPTIONS (#7806)

---
 test/test_datapoints.py              |  1 -
 torchvision/datapoints/_datapoint.py | 28 +++++++++++-----------------
 2 files changed, 11 insertions(+), 18 deletions(-)

diff --git a/test/test_datapoints.py b/test/test_datapoints.py
index 984caa2c345..4987b1991fc 100644
--- a/test/test_datapoints.py
+++ b/test/test_datapoints.py
@@ -209,4 +209,3 @@ def test_deepcopy(datapoint, requires_grad):
 
     assert type(datapoint_deepcopied) is type(datapoint)
     assert datapoint_deepcopied.requires_grad is requires_grad
-    assert datapoint_deepcopied.is_leaf
diff --git a/torchvision/datapoints/_datapoint.py b/torchvision/datapoints/_datapoint.py
index af6d5929d10..7032d518fe4 100644
--- a/torchvision/datapoints/_datapoint.py
+++ b/torchvision/datapoints/_datapoint.py
@@ -33,14 +33,9 @@ def _to_tensor(
     def wrap_like(cls: Type[D], other: D, tensor: torch.Tensor) -> D:
         return tensor.as_subclass(cls)
 
-    _NO_WRAPPING_EXCEPTIONS = {
-        torch.Tensor.clone: lambda cls, input, output: cls.wrap_like(input, output),
-        torch.Tensor.to: lambda cls, input, output: cls.wrap_like(input, output),
-        torch.Tensor.detach: lambda cls, input, output: cls.wrap_like(input, output),
-        # We don't need to wrap the output of `Tensor.requires_grad_`, since it is an inplace operation and thus
-        # retains the type automatically
-        torch.Tensor.requires_grad_: lambda cls, input, output: output,
-    }
+    # The ops in this set are those that should *preserve* the Datapoint type,
+    # i.e. they are exceptions to the "no wrapping" rule.
+    _NO_WRAPPING_EXCEPTIONS = {torch.Tensor.clone, torch.Tensor.to, torch.Tensor.detach, torch.Tensor.requires_grad_}
 
     @classmethod
     def __torch_function__(
@@ -76,22 +71,21 @@ def __torch_function__(
         with DisableTorchFunctionSubclass():
             output = func(*args, **kwargs or dict())
 
-            wrapper = cls._NO_WRAPPING_EXCEPTIONS.get(func)
-            # Apart from `func` needing to be an exception, we also require the primary operand, i.e. `args[0]`, to be
+        if func in cls._NO_WRAPPING_EXCEPTIONS and isinstance(args[0], cls):
+            # We also require the primary operand, i.e. `args[0]`, to be
             # an instance of the class that `__torch_function__` was invoked on. The __torch_function__ protocol will
             # invoke this method on *all* types involved in the computation by walking the MRO upwards. For example,
             # `torch.Tensor(...).to(datapoints.Image(...))` will invoke `datapoints.Image.__torch_function__` with
             # `args = (torch.Tensor(), datapoints.Image())` first. Without this guard, the original `torch.Tensor` would
             # be wrapped into a `datapoints.Image`.
-            if wrapper and isinstance(args[0], cls):
-                return wrapper(cls, args[0], output)
+            return cls.wrap_like(args[0], output)
 
-            # Inplace `func`'s, canonically identified with a trailing underscore in their name like `.add_(...)`,
-            # will retain the input type. Thus, we need to unwrap here.
-            if isinstance(output, cls):
-                return output.as_subclass(torch.Tensor)
+        if isinstance(output, cls):
+            # DisableTorchFunctionSubclass is ignored by inplace ops like `.add_(...)`,
+            # so for those, the output is still a Datapoint. Thus, we need to manually unwrap.
+            return output.as_subclass(torch.Tensor)
 
-            return output
+        return output
 
     def _make_repr(self, **kwargs: Any) -> str:
         # This is a poor man's implementation of the proposal in https://github.com/pytorch/pytorch/issues/76532.