Repository: pytorch/vision

Issue:
`Tensor.to` does not work with feature as reference
`torch.Tensor.to` accepts a tensor reference object as input in which case it converts to the reference dtype and device:

```py
>>> a = torch.zeros((), dtype=torch.float)
>>> b = torch.ones((), dtype=torch.int)
>>> b
tensor(1, dtype=torch.int32)
>>> b.to(a)
tensor(1.)
```

This also works for our custom features:

```py
>>> data = torch.rand(3, 2, 2)
>>> image = features.Image(data.to(torch.float64))
>>> image
Image([[[0.9595, 0.0947],
        [0.9553, 0.5563]],
       [[0.5435, 0.2975],
        [0.3037, 0.0863]],
       [[0.6253, 0.3481],
        [0.3518, 0.4499]]], dtype=torch.float64)
>>> image.to(torch.float32)
Image([[[0.9595, 0.0947],
        [0.9553, 0.5563]],
       [[0.5435, 0.2975],
        [0.3037, 0.0863]],
       [[0.6253, 0.3481],
        [0.3518, 0.4499]]])
>>> image.to(data)
Image([[[0.9595, 0.0947],
        [0.9553, 0.5563]],
       [[0.5435, 0.2975],
        [0.3037, 0.0863]],
       [[0.6253, 0.3481],
        [0.3518, 0.4499]]])

```

However, it doesn't work if we want to convert a plain tensor and use a custom feature as reference:

```py
>>> data.to(image)
AttributeError: 'Tensor' object has no attribute 'color_space'
```

cc @bjuncek @pmeier


Discussion:

Resolving patch:
diff --git a/torchvision/prototype/features/_feature.py b/torchvision/prototype/features/_feature.py
--- a/torchvision/prototype/features/_feature.py
+++ b/torchvision/prototype/features/_feature.py
@@ -89,6 +89,13 @@ def __torch_function__(
         with DisableTorchFunction():
             output = func(*args, **kwargs)
 
+        # 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(features.Image(...))` will invoke
+        # `features.Image.__torch_function__` first. The check below makes sure that we do not try to wrap in such a
+        # case.
+        if not isinstance(args[0], cls):
+            return output
+
         if func is torch.Tensor.clone:
             return cls.new_like(args[0], output)
         elif func is torch.Tensor.to:


Test patch:
diff --git a/test/test_prototype_features.py b/test/test_prototype_features.py
new file mode 100644
--- /dev/null
+++ b/test/test_prototype_features.py
@@ -0,0 +1,72 @@
+import torch
+from torchvision.prototype import features
+
+
+def test_isinstance():
+    assert isinstance(
+        features.Label([0, 1, 0], categories=["foo", "bar"]),
+        torch.Tensor,
+    )
+
+
+def test_wrapping_no_copy():
+    tensor = torch.tensor([0, 1, 0], dtype=torch.int64)
+    label = features.Label(tensor, categories=["foo", "bar"])
+
+    assert label.data_ptr() == tensor.data_ptr()
+
+
+def test_to_wrapping():
+    tensor = torch.tensor([0, 1, 0], dtype=torch.int64)
+    label = features.Label(tensor, categories=["foo", "bar"])
+
+    label_to = label.to(torch.int32)
+
+    assert type(label_to) is features.Label
+    assert label_to.dtype is torch.int32
+    assert label_to.categories is label.categories
+
+
+def test_to_feature_reference():
+    tensor = torch.tensor([0, 1, 0], dtype=torch.int64)
+    label = features.Label(tensor, categories=["foo", "bar"]).to(torch.int32)
+
+    tensor_to = tensor.to(label)
+
+    assert type(tensor_to) is torch.Tensor
+    assert tensor_to.dtype is torch.int32
+
+
+def test_clone_wrapping():
+    tensor = torch.tensor([0, 1, 0], dtype=torch.int64)
+    label = features.Label(tensor, categories=["foo", "bar"])
+
+    label_clone = label.clone()
+
+    assert type(label_clone) is features.Label
+    assert label_clone.data_ptr() != label.data_ptr()
+    assert label_clone.categories is label.categories
+
+
+def test_other_op_no_wrapping():
+    tensor = torch.tensor([0, 1, 0], dtype=torch.int64)
+    label = features.Label(tensor, categories=["foo", "bar"])
+
+    # any operation besides .to() and .clone() will do here
+    output = label * 2
+
+    assert type(output) is torch.Tensor
+
+
+def test_new_like():
+    tensor = torch.tensor([0, 1, 0], dtype=torch.int64)
+    label = features.Label(tensor, categories=["foo", "bar"])
+
+    # any operation besides .to() and .clone() will do here
+    output = label * 2
+
+    label_new = features.Label.new_like(label, output)
+
+    assert type(label_new) is features.Label
+    assert label_new.data_ptr() == output.data_ptr()
+    assert label_new.categories is label.categories