Repository: pytorch/vision

Issue:
Port `transforms.LinearTransformation` to `prototype.transforms`
cc @vfdev-5 @datumbox @bjuncek


Discussion:
@pmeier Out-of-curiosity, is any help needed for this feature or does it require too much knowledge/not advised for new contributors? :thinking: 
The removal of the "help wanted" label is explained in https://github.com/pytorch/vision/issues/5520#issuecomment-1061920594. In addition, https://github.com/pytorch/vision/issues/5336#issuecomment-1080586735 also applies. Thus, at this time, it is probably not a good idea to start. This has nothing to do with the knowledge / experience though. I'm happy to assign you as soon we have ironed out the problems.

If you are looking at ways to contribute, you might also have a look at [all issues tagged with the "help wanted" label](https://github.com/pytorch/vision/issues?q=is%3Aopen+is%3Aissue+label%3A%22help+wanted%22). That being said, please be aware that we (`torchvision` maintainers) are experiencing a heavy internal "load" at the moment that takes a lot of our time. I'm not saying this to discourage contributions, but rather as explanation why reviews might take some time.
Thanks for the details and the pointers @pmeier.
I think I found something to contribute in the meantime. 
Please let me know once things are a bit "lighter" on your side and you have more time. :ok_hand:  

Resolving patch:
diff --git a/torchvision/prototype/transforms/__init__.py b/torchvision/prototype/transforms/__init__.py
--- a/torchvision/prototype/transforms/__init__.py
+++ b/torchvision/prototype/transforms/__init__.py
@@ -37,7 +37,7 @@
     TenCrop,
 )
 from ._meta import ConvertBoundingBoxFormat, ConvertColorSpace, ConvertImageDtype
-from ._misc import GaussianBlur, Identity, Lambda, Normalize, ToDtype
+from ._misc import GaussianBlur, Identity, Lambda, LinearTransformation, Normalize, ToDtype
 from ._type_conversion import DecodeImage, LabelToOneHot, ToImagePIL, ToImageTensor
 
 from ._deprecated import Grayscale, RandomGrayscale, ToTensor, ToPILImage, PILToTensor  # usort: skip
diff --git a/torchvision/prototype/transforms/_misc.py b/torchvision/prototype/transforms/_misc.py
--- a/torchvision/prototype/transforms/_misc.py
+++ b/torchvision/prototype/transforms/_misc.py
@@ -1,7 +1,10 @@
 import functools
 from typing import Any, Callable, Dict, List, Sequence, Type, Union
 
+import PIL.Image
+
 import torch
+from torchvision.prototype import features
 from torchvision.prototype.transforms import functional as F, Transform
 from torchvision.transforms.transforms import _setup_size
 
@@ -32,6 +35,59 @@ def extra_repr(self) -> str:
         return ", ".join(extras)
 
 
+class LinearTransformation(Transform):
+    def __init__(self, transformation_matrix: torch.Tensor, mean_vector: torch.Tensor):
+        super().__init__()
+        if transformation_matrix.size(0) != transformation_matrix.size(1):
+            raise ValueError(
+                "transformation_matrix should be square. Got "
+                f"{tuple(transformation_matrix.size())} rectangular matrix."
+            )
+
+        if mean_vector.size(0) != transformation_matrix.size(0):
+            raise ValueError(
+                f"mean_vector should have the same length {mean_vector.size(0)}"
+                f" as any one of the dimensions of the transformation_matrix [{tuple(transformation_matrix.size())}]"
+            )
+
+        if transformation_matrix.device != mean_vector.device:
+            raise ValueError(
+                f"Input tensors should be on the same device. Got {transformation_matrix.device} and {mean_vector.device}"
+            )
+
+        self.transformation_matrix = transformation_matrix
+        self.mean_vector = mean_vector
+
+    def _transform(self, inpt: Any, params: Dict[str, Any]) -> Any:
+
+        if isinstance(inpt, features._Feature) and not isinstance(inpt, features.Image):
+            return inpt
+        elif isinstance(inpt, PIL.Image.Image):
+            raise TypeError("Unsupported input type")
+
+        # Image instance after linear transformation is not Image anymore due to unknown data range
+        # Thus we will return Tensor for input Image
+
+        shape = inpt.shape
+        n = shape[-3] * shape[-2] * shape[-1]
+        if n != self.transformation_matrix.shape[0]:
+            raise ValueError(
+                "Input tensor and transformation matrix have incompatible shape."
+                + f"[{shape[-3]} x {shape[-2]} x {shape[-1]}] != "
+                + f"{self.transformation_matrix.shape[0]}"
+            )
+
+        if inpt.device.type != self.mean_vector.device.type:
+            raise ValueError(
+                "Input tensor should be on the same device as transformation matrix and mean vector. "
+                f"Got {inpt.device} vs {self.mean_vector.device}"
+            )
+
+        flat_tensor = inpt.view(-1, n) - self.mean_vector
+        transformed_tensor = torch.mm(flat_tensor, self.transformation_matrix)
+        return transformed_tensor.view(shape)
+
+
 class Normalize(Transform):
     def __init__(self, mean: List[float], std: List[float]):
         super().__init__()


Test patch:
diff --git a/test/test_prototype_transforms.py b/test/test_prototype_transforms.py
--- a/test/test_prototype_transforms.py
+++ b/test/test_prototype_transforms.py
@@ -1487,3 +1487,36 @@ def test__transform_bounding_box_clamping(self, mocker):
         transform(bounding_box)
 
         mock.assert_called_once()
+
+
+class TestLinearTransformation:
+    def test_assertions(self):
+        with pytest.raises(ValueError, match="transformation_matrix should be square"):
+            transforms.LinearTransformation(torch.rand(2, 3), torch.rand(5))
+
+        with pytest.raises(ValueError, match="mean_vector should have the same length"):
+            transforms.LinearTransformation(torch.rand(3, 3), torch.rand(5))
+
+    @pytest.mark.parametrize(
+        "inpt",
+        [
+            122 * torch.ones(1, 3, 8, 8),
+            122.0 * torch.ones(1, 3, 8, 8),
+            features.Image(122 * torch.ones(1, 3, 8, 8)),
+            PIL.Image.new("RGB", (8, 8), (122, 122, 122)),
+        ],
+    )
+    def test__transform(self, inpt):
+
+        v = 121 * torch.ones(3 * 8 * 8)
+        m = torch.ones(3 * 8 * 8, 3 * 8 * 8)
+        transform = transforms.LinearTransformation(m, v)
+
+        if isinstance(inpt, PIL.Image.Image):
+            with pytest.raises(TypeError, match="Unsupported input type"):
+                transform(inpt)
+        else:
+            output = transform(inpt)
+            assert isinstance(output, torch.Tensor)
+            assert output.unique() == 3 * 8 * 8
+            assert output.dtype == inpt.dtype