Update

[ghstack-poisoned]

Author: jainapurva

test/test_utils.py

@@ -41,9 +41,11 @@ def __init__(self, data):
                 self.data = data
 
         l = torch.nn.Linear(10, 10)
-        with self.assertRaisesRegex(NotImplementedError, "arg_types"):
+        with self.assertRaisesRegex(NotImplementedError, "Subclasses must implement"):
             l.weight = torch.nn.Parameter(MyTensor(l.weight))
 
+        assert MyTensor(l.weight).get_layout() is None
+
 
 if __name__ == "__main__":
     unittest.main()

torchao/utils.py

@@ -9,6 +9,7 @@
 
 import torch
 import torch.nn.utils.parametrize as parametrize
+from torch.utils._python_dispatch import return_and_correct_aliasing
 
 __all__ = [
     "benchmark_model",
@@ -512,9 +513,36 @@ def _get_tensor_impl_constructor(
     return tensor_class._LAYOUT_CONSTRUCTOR_TABLE[layout_class]
 
 
+def _get_to_kwargs(self, *args, **kwargs):
+    # `torch._C._nn._parse_to` can't handle `layout` argument
+    for arg in args:
+        if isinstance(arg, torch.layout):
+            args.remove(arg)
+    if "layout" in kwargs:
+        kwargs.pop("layout")
+    # ignoring `non_blocking` and `memory_format` args since these are not
+    # very useful for most of the tensor subclasses
+    # if in the future there are use cases that need these, we'd recommend
+    # to override `_get_to_kwargs` and return these args
+    device, dtype, _, _ = torch._C._nn._parse_to(*args, **kwargs)
+    device = self.device if device is None else device
+    dtype = self.dtype if dtype is None else dtype
+    kwargs = {
+        "device": device,
+        "dtype": dtype,
+    }
+    return kwargs
+
+
+"""
+TensorAOBase subclass is a util tensor subclass that provides commonly used functions, and should be inherited to define a new tensor subclass
+"""
+
+
 class TorchAOBaseTensor(torch.Tensor):
     """A util tensor subclass that provides commonly used functions
-       new tensor subclass can inherit it to get all the utility functions
+       new tensor subclass can inherit it to get all the utility functions, and
+       should be inherited to define a new tensor subclass
 
        class MyTensor(TorchAOBaseTensor):
            pass
@@ -539,39 +567,45 @@ def _(func, types, args, kwargs):
             class PlainAQTTensorImpl(...):
                 ...
 
-         `get_tensor_impl_constructor`:
+        `get_tensor_impl_constructor`:
             get_tensor_impl_constructor = MyTensor.get_tensor_impl_constructor
             # in constructor of MyTensor:
             tensor_impl_ctr = get_tensor_impl_constructor(type(_layout))
             tensor_impl = tensor_impl_ctr(data, scale, zero_point, _layout)
 
+        `__tensor_flatten__` and `__tensor_unflatten__`: Used for tensor serialization/deserialization, must be defined in the subclass
+
+        `__repr__`: Used for tensor representation, must be defined in the subclass
+
+        `_apply_fn_to_data`: Used for applying a function to the data of the tensor, must be defined in the subclass
     """
 
     implements = classmethod(_implements)
     __torch_dispatch__ = classmethod(_dispatch__torch_dispatch__)
     __torch_function__ = classmethod(_dispatch__torch_function__)
     register_layout = classmethod(_register_layout)
     get_tensor_impl_constructor = classmethod(_get_tensor_impl_constructor)
+    _get_to_kwargs = _get_to_kwargs
+
+    def __tensor_flatten__(self):
+        raise NotImplementedError("Subclasses must implement __tensor_flatten__")
+
+    @classmethod
+    def __tensor_unflatten__(
+        cls, tensor_data_dict, tensor_attributes, outer_size, outer_stride
+    ):
+        raise NotImplementedError("Subclasses must implement __tensor_unflatten__")
+
+    def __repr__(self):
+        raise NotImplementedError("Subclasses must implement __repr__")
+
+    def _apply_fn_to_data(self, fn):
+        raise NotImplementedError("Subclasses must implement _apply_fn_to_data")
 
-    def _get_to_kwargs(self, *args, **kwargs):
-        # `torch._C._nn._parse_to` can't handle `layout` argument
-        for arg in args:
-            if isinstance(arg, torch.layout):
-                args.remove(arg)
-        if "layout" in kwargs:
-            kwargs.pop("layout")
-        # ignoring `non_blocking` and `memory_format` args since these are not
-        # very useful for most of the tensor subclasses
-        # if in the future there are use cases that need these, we'd recommend
-        # to override `_get_to_kwargs` and return these args
-        device, dtype, _, _ = torch._C._nn._parse_to(*args, **kwargs)
-        device = self.device if device is None else device
-        dtype = self.dtype if dtype is None else dtype
-        kwargs = {
-            "device": device,
-            "dtype": dtype,
-        }
-        return kwargs
+    def get_layout(self):
+        if not hasattr(self, "_layout"):
+            return None
+        return self._layout
 
 
 def fill_defaults(args, n, defaults_tail):
@@ -599,7 +633,7 @@ def fill_defaults(args, n, defaults_tail):
     return r
 
 
-## Deprecated, will be deleted in the future
+# Deprecated, will be deleted in the future
 def _torch_version_at_least(min_version):
     return is_fbcode() or version("torch") >= min_version