Support more ops in TorchAOBaseTensor (#2609)

Summary:
* detach
* clone
* alias
* contiguous
* copy_
* to

Test Plan:
python test/test_utils.py

Reviewers:

Subscribers:

Tasks:

Tags:

stack-info: PR: #2598, branch: jerryzh168/stack/13

Author: jerryzh168

test/test_utils.py

@@ -7,8 +7,8 @@
 from unittest.mock import patch
 
 import torch
-from torch.utils._python_dispatch import return_and_correct_aliasing
 
+from torchao.testing.utils import skip_if_no_cuda
 from torchao.utils import TorchAOBaseTensor, torch_version_at_least
 
 
@@ -47,37 +47,39 @@ def __init__(self, data):
                 self.data = data
 
         l = torch.nn.Linear(10, 10)
+        # since we did not define `tensor_data_names` and `tensor_attribute_names` for MyTensor
+        # the following call will error out because `detach` is defined in `TorchAOBaseTensor`
+        # but would rely on `tensor_data_names` and `tensor_attribute_names` being defined for it to work
+        # user could either specify `tensor_data_names` and `tensor_attribute_names` or manually implement
+        # detach op
         with self.assertRaisesRegex(NotImplementedError, "arg_types"):
             l.weight = torch.nn.Parameter(MyTensor(l.weight))
 
+    @skip_if_no_cuda()
     def test_default_impls(self):
         """Making sure some common functions has default implementations, such as
         __tensor_unflatten__, __tensor_flatten__, _apply_fn_to_data, __repr__, to
         """
 
         class MyTensor(TorchAOBaseTensor):
             tensor_data_names = ["qdata"]
-            tensor_attribute_names = ["attr"]
+            tensor_attribute_names = ["attr", "device"]
 
-            def __new__(cls, qdata, attr):
+            def __new__(cls, qdata, attr, device=None):
                 shape = qdata.shape
-                return torch.Tensor._make_wrapper_subclass(cls, shape)  # type: ignore[attr-defined]
+                if device is None:
+                    device = qdata.device
+                kwargs = {"device": device}
+                return torch.Tensor._make_wrapper_subclass(cls, shape, **kwargs)  # type: ignore[attr-defined]
 
-            def __init__(self, qdata, attr):
+            def __init__(self, qdata, attr, device=None):
                 self.qdata = qdata
                 self.attr = attr
 
-        implements = MyTensor.implements
-
-        @implements(torch.ops.aten.detach.default)
-        def _(func, types, args, kwargs):
-            return return_and_correct_aliasing(
-                func, args, kwargs, args[0]._apply_fn_to_data(torch.detach)
-            )
-
         l = torch.nn.Linear(1, 1)
         l.weight = torch.nn.Parameter(MyTensor(l.weight, "attr"))
         lp_tensor = l.weight
+        # test __tensor_flatten__ and __tensor_unflatten__
         tensor_data_name_dict, tensor_attributes = lp_tensor.__tensor_flatten__()
         tensor_data_dict = {
             name: getattr(lp_tensor, name) for name in tensor_data_name_dict
@@ -89,8 +91,35 @@ def _(func, types, args, kwargs):
         )
         self.assertTrue(torch.equal(lp_tensor.qdata, reconstructed.qdata))
         self.assertEqual(lp_tensor.attr, reconstructed.attr)
+
+        # `to` / `_to_copy`
+        original_device = lp_tensor.device
+        lp_tensor = lp_tensor.to("cuda")
+        self.assertEqual(lp_tensor.device.type, "cuda")
+        lp_tensor = lp_tensor.to(original_device)
+        self.assertEqual(lp_tensor.device, original_device)
+
+        # __repr__
         print(lp_tensor)
 
+        # other ops
+        lp_tensor = lp_tensor.detach()
+        # explicitly testing aten.alias
+        lp_tensor = torch.ops.aten.alias(lp_tensor)
+        lp_tensor = lp_tensor.clone()
+        lp_tensor = lp_tensor.contiguous()
+
+        # copy_
+        another_tensor = torch.nn.Linear(1, 1).weight
+        # attribute has to be the same
+        another_lp_tensor = MyTensor(another_tensor, "attr")
+        # initially tensor values are not the same
+        self.assertNotEqual(lp_tensor.qdata[0], another_lp_tensor.qdata[0])
+        lp_tensor.copy_(another_lp_tensor)
+        self.assertEqual(lp_tensor.attr, "attr")
+        # after copy_, the tensor values should match
+        self.assertEqual(lp_tensor.qdata[0], another_lp_tensor.qdata[0])
+
 
 if __name__ == "__main__":
     unittest.main()

torchao/utils.py

@@ -15,6 +15,7 @@
 
 import torch
 import torch.nn.utils.parametrize as parametrize
+from torch.utils._python_dispatch import return_and_correct_aliasing
 
 __all__ = [
     "benchmark_model",
@@ -409,6 +410,9 @@ def _(func, types, args, kwargs):
     if not hasattr(cls, "_ATEN_OP_OR_TORCH_FN_TABLE"):
         cls._ATEN_OP_OR_TORCH_FN_TABLE = {}
 
+    if cls not in cls._ATEN_OP_OR_TORCH_FN_TABLE:
+        cls._ATEN_OP_OR_TORCH_FN_TABLE[cls] = {}
+
     if not isinstance(aten_ops_or_torch_fns, (list, tuple)):
         aten_ops_or_torch_fns = [aten_ops_or_torch_fns]
 
@@ -419,12 +423,83 @@ def decorator(func):
             def wrapper(f, types, args, kwargs):
                 return func(f, types, args, kwargs)
 
-            cls._ATEN_OP_OR_TORCH_FN_TABLE[op] = wrapper
+            cls._ATEN_OP_OR_TORCH_FN_TABLE[cls][op] = wrapper
         return func
 
     return decorator
 
 
+def _implements_common_tensor_ops(cls):
+    implements = cls.implements
+    aten = torch.ops.aten
+
+    @implements(
+        [aten.detach.default, aten.clone.default, aten.alias.default, aten.contiguous]
+    )
+    def _(func, types, args, kwargs):
+        return return_and_correct_aliasing(
+            func,
+            args,
+            kwargs,
+            args[0]._apply_fn_to_data(lambda x: func(x, *args[1:], **kwargs)),
+        )
+
+    def _same_metadata(self: TorchAOBaseTensor, src: TorchAOBaseTensor) -> bool:
+        _tensor_shape_match = all(
+            getattr(self, t_name).shape == getattr(src, t_name).shape
+            for t_name in self.tensor_data_names
+        )
+        _attr_match = all(
+            getattr(self, a_name) == getattr(src, a_name)
+            for a_name in self.tensor_attribute_names
+        )
+        return (
+            type(self) == type(src)
+            and self.shape == src.shape
+            and _tensor_shape_match
+            and _attr_match
+        )
+
+    @implements(aten.copy_.default)
+    def _(func, types, args, kwargs):
+        self = args[0]
+        src = args[1]
+        if _same_metadata(self, src):
+            self_tensors = self.__tensor_flatten__()[0]
+            for tensor_name in self_tensors:
+                getattr(self, tensor_name).copy_(getattr(src, tensor_name))
+            return
+        raise ValueError(
+            f"Not supported args for copy_ due to metadata mismatch: {args[0], args[1]}"
+        )
+
+    @implements(aten._to_copy.default)
+    def _(func, types, args, kwargs):
+        self = args[0]
+        if hasattr(self, "tensor_data_names") and hasattr(
+            self, "tensor_attribute_names"
+        ):
+            kwargs = self._get_to_kwargs(*args[1:], **kwargs)
+            device = kwargs.pop("device")
+            tensors = [
+                getattr(self, name).to(device) for name in self.tensor_data_names
+            ]
+            # change device
+            tensor_attributes = [
+                getattr(self, attr_name) if attr_name != "device" else device
+                for attr_name in self.tensor_attribute_names
+            ]
+            t = self.__class__(
+                *tensors,
+                *tensor_attributes,
+            )
+            return return_and_correct_aliasing(func, args, kwargs, t)
+
+        raise NotImplementedError(
+            "Subclasses must implement `aten._to_copy.default` or specify `tensor_data_names` and `tensor_attribute_names` for tensor class or tensor instance before using it"
+        )
+
+
 def _dispatch__torch_function__(cls, func, types, args=(), kwargs=None):
     """Use this util function for a common `__torch_function__` implementation
     that dispatches to ops/functions registered with `_implements`
@@ -436,9 +511,10 @@ class MyTensor(torch.Tensor):
     kwargs = {} if kwargs is None else kwargs
     if (
         hasattr(cls, "_ATEN_OP_OR_TORCH_FN_TABLE")
-        and func in cls._ATEN_OP_OR_TORCH_FN_TABLE
+        and cls in cls._ATEN_OP_OR_TORCH_FN_TABLE
+        and func in cls._ATEN_OP_OR_TORCH_FN_TABLE[cls]
     ):
-        return cls._ATEN_OP_OR_TORCH_FN_TABLE[func](func, types, args, kwargs)
+        return cls._ATEN_OP_OR_TORCH_FN_TABLE[cls][func](func, types, args, kwargs)
 
     with torch._C.DisableTorchFunctionSubclass():
         return func(*args, **kwargs)
@@ -454,9 +530,10 @@ class MyTensor(torch.Tensor):
     """
     if (
         hasattr(cls, "_ATEN_OP_OR_TORCH_FN_TABLE")
-        and func in cls._ATEN_OP_OR_TORCH_FN_TABLE
+        and cls in cls._ATEN_OP_OR_TORCH_FN_TABLE
+        and func in cls._ATEN_OP_OR_TORCH_FN_TABLE[cls]
     ):
-        return cls._ATEN_OP_OR_TORCH_FN_TABLE[func](func, types, args, kwargs)
+        return cls._ATEN_OP_OR_TORCH_FN_TABLE[cls][func](func, types, args, kwargs)
 
     arg_types = tuple(type(arg) for arg in args)
     kwarg_types = {k: type(arg) for k, arg in kwargs.items()}
@@ -576,7 +653,28 @@ class PlainAQTTensorImpl(...):
 
     """
 
+    @classmethod
+    def __init_subclass__(cls, **kwargs):
+        if not hasattr(cls, "_ATEN_OP_OR_TORCH_FN_TABLE"):
+            cls._ATEN_OP_OR_TORCH_FN_TABLE = {}
+
+        if cls not in cls._ATEN_OP_OR_TORCH_FN_TABLE:
+            cls._ATEN_OP_OR_TORCH_FN_TABLE[cls] = {}
+
+        # define the common ops if the tensor_data_names and tensor_attribute_names are defined
+        if hasattr(cls, "tensor_data_names") and hasattr(cls, "tensor_attribute_names"):
+            cls._implements_common_tensor_ops()
+
+        # inherit the torch function and dispatch implementations from direct parent classes
+        # e.g. for `class C(B, A)`, C.__bases__ == (B, A)
+        for parent in cls.__bases__:
+            if parent in cls._ATEN_OP_OR_TORCH_FN_TABLE:
+                cls._ATEN_OP_OR_TORCH_FN_TABLE[cls].update(
+                    cls._ATEN_OP_OR_TORCH_FN_TABLE[parent]
+                )
+
     implements = classmethod(_implements)
+    _implements_common_tensor_ops = classmethod(_implements_common_tensor_ops)
     __torch_dispatch__ = classmethod(_dispatch__torch_dispatch__)
     __torch_function__ = classmethod(_dispatch__torch_function__)
     register_layout = classmethod(_register_layout)
@@ -591,7 +689,7 @@ def __tensor_flatten__(self):
                 getattr(self, attr) for attr in self.tensor_attribute_names
             ]
         raise NotImplementedError(
-            "Subclasses must implement __tensor_flatten__ or specify `tensor_data_names` and `tensor_attribute_names` for tensor class or tensor instance"
+            "Subclasses should implement __tensor_flatten__ or specify `tensor_data_names` and `tensor_attribute_names` for tensor class or tensor instance before using it"
         )
 
     @classmethod
@@ -602,13 +700,20 @@ def __tensor_unflatten__(
         return cls(*tensors, *tensor_attributes)
 
     def _apply_fn_to_data(self, fn):
-        tensors = [fn(getattr(self, attr)) for attr in self.tensor_data_names]
-        tensor_attributes = [
-            getattr(self, attr) for attr in self.tensor_attribute_names
-        ]
-        return self.__class__(
-            *tensors,
-            *tensor_attributes,
+        if hasattr(self, "tensor_data_names") and hasattr(
+            self, "tensor_attribute_names"
+        ):
+            tensors = [fn(getattr(self, attr)) for attr in self.tensor_data_names]
+            tensor_attributes = [
+                getattr(self, attr) for attr in self.tensor_attribute_names
+            ]
+            return self.__class__(
+                *tensors,
+                *tensor_attributes,
+            )
+
+        raise NotImplementedError(
+            "Subclasses should implement _apply_fn_to_data or specify `tensor_data_names` and `tensor_attribute_names` for tensor class or tensor instance before using it"
         )
 
     def __repr__(self):
@@ -624,7 +729,10 @@ def __repr__(self):
                     f", {tensor_attribute_name}={getattr(self, tensor_attribute_name)}"
                 )
             return f"{self.__class__.__name__}({repr_str})"
-        raise NotImplementedError("Subclasses must implement __repr__")
+
+        raise NotImplementedError(
+            "Subclasses must implement __repr__ or specify `tensor_data_names` and `tensor_attribute_names` for tensor class or tensor instance before using it"
+        )
 
     def get_layout(self):
         if not hasattr(self, "_layout"):