addressing comments

Author: liangel-02

test/prototype/safetensors/test_safetensors.py renamed to test/prototype/safetensors/test_safetensors_support.py

@@ -8,7 +8,7 @@
 )
 
 from torchao import quantize_
-from torchao.prototype.quantization.safetensors_support import (
+from torchao.prototype.safetensors.safetensors_support import (
     load_tensor_state_dict,
     save_tensor_state_dict,
 )

torchao/prototype/quantization/safetensors_serialization.py renamed to torchao/prototype/safetensors/safetensors_serialization.py

@@ -5,14 +5,18 @@
 from typing import Any, Dict
 
 import torch
+import torchao
 
 from torchao.quantization import Float8Tensor
-
-ALLOWED_AO_MODULES = {
-    "torchao.quantization",
-    "torchao.dtypes",
-    "torchao.quantization.quantize_.common",
-    "torchao.quantization.quantize_.workflows",
+from torchao.quantization.quantize_.workflows import QuantizeTensorToFloat8Kwargs
+from torchao.quantization.quantize_.common import KernelPreference
+
+ALLOWED_CLASSES = {
+    "Float8Tensor": torchao.quantization.Float8Tensor,
+    "Float8MMConfig": torchao.float8.inference.Float8MMConfig,
+    "QuantizeTensorToFloat8Kwargs": QuantizeTensorToFloat8Kwargs,
+    "PerRow": torchao.quantization.granularity.PerRow,
+    "KernelPreference": KernelPreference,
 }
 
 
@@ -58,7 +62,7 @@ def default(self, o):
             return {"_type": "torch.dtype", "_data": str(o).split(".")[-1]}
 
         if isinstance(o, enum.Enum):
-            # Store the full path for enums to ensure uniqueness
+            # Store the full class name for enums to ensure uniqueness
             return {"_type": f"{o.__class__.__name__}", "_data": o.name}
 
         if isinstance(o, list):
@@ -81,6 +85,8 @@ def encode_value(self, value):
         except TypeError:
             pass
 
+        # Default case - return as is
+        # (This will be processed by standard JSON encoder later)
         return value
 
 
@@ -97,19 +103,11 @@ def object_from_dict(data: Dict[str, Any]):
     if type_path == "torch.dtype":
         return getattr(torch, obj_data)
 
-    # Try to find the class in any of the allowed modules
-    cls = None
-    for module_path in ALLOWED_AO_MODULES:
-        try:
-            module = importlib.import_module(module_path)
-            cls = getattr(module, type_path)
-            break  # Found the class, exit the loop
-        except (ImportError, AttributeError):
-            continue  # Try the next module
+    cls = ALLOWED_CLASSES.get(type_path)
 
     # If we couldn't find the class in any allowed module, raise an error
     if cls is None:
-        allowed_modules_str = ", ".join(ALLOWED_AO_MODULES)
+        allowed_modules_str = ", ".join(ALLOWED_CLASSES)
         raise ValueError(
             f"Failed to find class {type_path} in any of the allowed modules: {allowed_modules_str}"
         )

torchao/prototype/quantization/safetensors_support.py renamed to torchao/prototype/safetensors/safetensors_support.py

@@ -5,7 +5,7 @@
 import torch
 from safetensors.torch import load_file, save_file
 
-from torchao.prototype.quantization.safetensors_serialization import (
+from torchao.prototype.safetensors.safetensors_serialization import (
     Float8TensorAttributeJSONEncoder,
     object_from_dict,
 )
@@ -118,7 +118,7 @@ def save_tensor_state_dict(
                 tensors_dict[tensor_data_name] = getattr(tensor, tensor_data_name)
 
             metadata = json.dumps(tensor, cls=Float8TensorAttributeJSONEncoder)
-        elif isinstance(tensor, torch.Tensor):
+        elif type(tensor) is torch.Tensor:
             tensors_dict = {"_data": tensor}
             metadata = json.dumps({"_type": torch.Tensor.__name__})
         else: