[test only] testing adding optioanl tensor arg to float8 tensor

Summary:

Test Plan:

Reviewers:

Subscribers:

Tasks:

Tags:

stack-info: PR: #2840, branch: jerryzh168/stack/33

Author: jerryzh168

torchao/quantization/quantize_/workflows/float8/float8_tensor.py

@@ -95,6 +95,7 @@ class Float8Tensor(TorchAOBaseTensor):
 
     tensor_data_names = ["qdata", "scale"]
     tensor_attribute_names = []
+    optional_tensor_data_names = ["test_only_data"]
     optional_tensor_attribute_names = [
         "block_size",
         "mm_config",
@@ -109,6 +110,7 @@ def __new__(
         cls,
         qdata: torch.Tensor,
         scale: torch.Tensor,
+        test_only_data: Optional[torch.Tensor] = None,
         block_size: Optional[List[int]] = None,
         mm_config: Optional[Float8MMConfig] = None,
         hp_value_lb: Optional[float] = None,
@@ -128,6 +130,7 @@ def __init__(
         self,
         qdata: torch.Tensor,
         scale: torch.Tensor,
+        test_only_data: Optional[torch.Tensor] = None,
         block_size: Optional[List[int]] = None,
         mm_config: Optional[Float8MMConfig] = None,
         hp_value_lb: Optional[float] = None,
@@ -138,6 +141,7 @@ def __init__(
     ):
         self.qdata = qdata
         self.scale = scale
+        self.test_only_data = test_only_data
         self.block_size = block_size
         self.mm_config = mm_config
         self.hp_value_lb = hp_value_lb
@@ -152,6 +156,11 @@ def __repr__(self):
             f"{self.shape=}, {self.device=}, {self.dtype=})"
         )
 
+    def __setstate__(self, state):
+        torch._utils._set_obj_state(self, state)
+        if "test_only_data" not in self.__dict__:
+            self.test_only_data = None
+
     def _quantization_type(self):
         return f"{self.act_quant_kwargs=}, {self.block_size=}, {self.mm_config=}, {self.scale.shape=}"