Add _apply_fn_to_data method to TorchAOBaseTensor base class

   - Implements generic pattern for applying functions to tensor components
   - Uses __tensor_flatten__ and __tensor_unflatten__ pattern
   - Fixes pytorch#2349

Author: krikera

test/test_utils.py

@@ -49,6 +49,14 @@ def __init__(self, data):
         with self.assertRaisesRegex(NotImplementedError, "arg_types"):
             l.weight = torch.nn.Parameter(MyTensor(l.weight))
 
+    def test_apply_fn_to_data(self):
+        self.assertTrue(hasattr(TorchAOBaseTensor, "_apply_fn_to_data"))
+        self.assertTrue(callable(getattr(TorchAOBaseTensor, "_apply_fn_to_data")))
+
+        method = getattr(TorchAOBaseTensor, "_apply_fn_to_data")
+        self.assertFalse(isinstance(method, classmethod))
+        self.assertFalse(isinstance(method, staticmethod))
+
 
 if __name__ == "__main__":
     unittest.main()

torchao/utils.py

@@ -577,6 +577,19 @@ def __tensor_unflatten__(
     ):
         raise NotImplementedError("Subclasses must implement __tensor_unflatten__")
 
+    def _apply_fn_to_data(self, fn: Callable):
+        """Applies a fn to all tensor components stored on this class"""
+        tensor_names, ctx = self.__tensor_flatten__()
+        new_tensors = {}
+        for name in tensor_names:
+            new_tensors[name] = fn(getattr(self, name))
+        return self.__class__.__tensor_unflatten__(
+            new_tensors,
+            ctx,
+            None,
+            None,
+        )
+
     def __repr__(self):
         raise NotImplementedError("Subclasses must implement __repr__")