Add more utils to TorchAOBaseTensor (#2597)

Summary:
Added default impls for:
* __tensor_flatten__ and __tensor_unflatten__ when tensor_data_names and tensor_attribute_names are defined
* __repr__
* _apply_fn_to_data

Next
* more op definitions

Test Plan:
python test/test_utils.py

Reviewers:

Subscribers:

Tasks:

Tags:

stack-info: PR: #2597, branch: jerryzh168/stack/12

Author: jerryzh168

test/test_utils.py

@@ -7,6 +7,7 @@
 from unittest.mock import patch
 
 import torch
+from torch.utils._python_dispatch import return_and_correct_aliasing
 
 from torchao.utils import TorchAOBaseTensor, torch_version_at_least
 
@@ -49,6 +50,47 @@ def __init__(self, data):
         with self.assertRaisesRegex(NotImplementedError, "arg_types"):
             l.weight = torch.nn.Parameter(MyTensor(l.weight))
 
+    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"]
+
+            def __new__(cls, qdata, attr):
+                shape = qdata.shape
+                return torch.Tensor._make_wrapper_subclass(cls, shape)  # type: ignore[attr-defined]
+
+            def __init__(self, qdata, attr):
+                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
+        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
+        }
+        outer_size = lp_tensor.size()
+        outer_stride = lp_tensor.stride()
+        reconstructed = type(lp_tensor).__tensor_unflatten__(
+            tensor_data_dict, tensor_attributes, outer_size, outer_stride
+        )
+        self.assertTrue(torch.equal(lp_tensor.qdata, reconstructed.qdata))
+        self.assertEqual(lp_tensor.attr, reconstructed.attr)
+        print(lp_tensor)
+
 
 if __name__ == "__main__":
     unittest.main()

torchao/utils.py

@@ -584,15 +584,46 @@ class PlainAQTTensorImpl(...):
     _get_to_kwargs = _get_to_kwargs
 
     def __tensor_flatten__(self):
-        raise NotImplementedError("Subclasses must implement __tensor_flatten__")
+        if hasattr(self, "tensor_data_names") and hasattr(
+            self, "tensor_attribute_names"
+        ):
+            return self.tensor_data_names, [
+                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"
+        )
 
     @classmethod
     def __tensor_unflatten__(
         cls, tensor_data_dict, tensor_attributes, outer_size, outer_stride
     ):
-        raise NotImplementedError("Subclasses must implement __tensor_unflatten__")
+        tensors = [tensor_data_dict[name] for name in cls.tensor_data_names]
+        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,
+        )
 
     def __repr__(self):
+        if hasattr(self, "tensor_data_names") and hasattr(
+            self, "tensor_attribute_names"
+        ):
+            repr_str = ""
+            repr_str += f"{self.tensor_data_names[0]}={getattr(self, self.tensor_data_names[0])}"
+            for tensor_data_name in self.tensor_data_names[1:]:
+                repr_str += f", {tensor_data_name}={getattr(self, tensor_data_name)}"
+            for tensor_attribute_name in self.tensor_attribute_names:
+                repr_str += (
+                    f", {tensor_attribute_name}={getattr(self, tensor_attribute_name)}"
+                )
+            return f"{self.__class__.__name__}({repr_str})"
         raise NotImplementedError("Subclasses must implement __repr__")
 
     def get_layout(self):