Add repr method on tensor subclass

Differential Revision: D58786586 Pull Request resolved: #397
pytorch · Jun 19, 2024 · ea71549 · ea71549
1 parent f5b6ec9
commit ea71549
Showing 1 changed file with 4 additions and 0 deletions.
diff --git a/torchao/dtypes/aqt.py b/torchao/dtypes/aqt.py
@@ -480,6 +480,10 @@ def _apply_fn_to_data(self, fn):
         self.scale_and_zero = fn(self.scale_and_zero)
         return self
 
+    def __repr__(self):
+        int_data, scale, zero_point = self.get_plain()
+        return f"TensorCoreTiledAQTLayout(int_data={int_data}, scale={scale}, zero_point={zero_point})"
+
     @classmethod
     def __torch_dispatch__(cls, func, types, args, kwargs):
         kwargs = {} if kwargs is None else kwargs