Move floatx_tensor_core_layout to prototype/dtypes (#3317)

Author: jainapurva

benchmarks/benchmark_fp6.py

@@ -9,7 +9,7 @@
 from tqdm import tqdm
 
 from torchao.dtypes import to_affine_quantized_fpx
-from torchao.dtypes.floatx import FloatxTensorCoreLayout
+from torchao.prototype.dtypes.floatx import FloatxTensorCoreLayout
 from torchao.utils import benchmark_torch_function_in_microseconds
 
 

docs/source/api_ref_dtypes.rst

@@ -20,7 +20,6 @@ Layouts and Tensor Subclasses
     TensorCoreTiledLayout
     Float8Layout
     FloatxTensor
-    FloatxTensorCoreLayout
     MarlinSparseLayout
     Int4CPULayout
     CutlassSemiSparseLayout
@@ -52,6 +51,7 @@ Prototype
     Int8DynamicActInt4WeightCPULayout
     MarlinQQQTensor
     MarlinQQQLayout
+    FloatxTensorCoreLayout
     UintxLayout
 
 ..

test/dtypes/test_floatx.py

@@ -14,20 +14,20 @@
     run_tests,
 )
 
-from torchao.dtypes.floatx import (
+from torchao.prototype.custom_fp_utils import (
+    _f32_to_floatx_unpacked,
+    _floatx_unpacked_to_f32,
+)
+from torchao.prototype.dtypes.floatx import (
     FloatxTensorCoreLayout,
     from_scaled_tc_floatx,
     to_scaled_tc_floatx,
 )
-from torchao.dtypes.floatx.floatx_tensor_core_layout import (
+from torchao.prototype.dtypes.floatx.floatx_tensor_core_layout import (
     FloatxTensorCoreAQTTensorImpl,
     _pack_tc_floatx,
     _pack_tc_fp6,
 )
-from torchao.prototype.custom_fp_utils import (
-    _f32_to_floatx_unpacked,
-    _floatx_unpacked_to_f32,
-)
 from torchao.quantization import (
     FPXWeightOnlyConfig,
     quantize_,

torchao/dtypes/affine_quantized_tensor.py

@@ -136,7 +136,8 @@ def dequantize(self, output_dtype: Optional[torch.dtype] = None) -> torch.Tensor
         if output_dtype is None:
             output_dtype = self.dtype
 
-        from torchao.dtypes.floatx import Float8Layout, FloatxTensorCoreLayout
+        from torchao.dtypes.floatx import Float8Layout
+        from torchao.prototype.dtypes.floatx import FloatxTensorCoreLayout
 
         if isinstance(self._layout, FloatxTensorCoreLayout):
             int_data, scale = self.tensor_impl.get_plain()
@@ -539,7 +540,7 @@ def from_hp_to_fpx(
         _layout: Layout,
     ):
         """Create a floatx AffineQuantizedTensor from a high precision tensor. Floatx is represented as ebits and mbits, and supports the representation of float1-float7."""
-        from torchao.dtypes.floatx import FloatxTensorCoreLayout
+        from torchao.prototype.dtypes.floatx import FloatxTensorCoreLayout
 
         assert isinstance(_layout, FloatxTensorCoreLayout), (
             f"Only FloatxTensorCoreLayout is supported for floatx, got {_layout}"

torchao/dtypes/affine_quantized_tensor_ops.py

@@ -21,10 +21,6 @@
     _linear_fp_act_fp8_weight_check,
     _linear_fp_act_fp8_weight_impl,
 )
-from torchao.dtypes.floatx.floatx_tensor_core_layout import (
-    _linear_f16_bf16_act_floatx_weight_check,
-    _linear_f16_bf16_act_floatx_weight_impl,
-)
 from torchao.dtypes.uintx.int4_cpu_layout import (
     _linear_fp_act_uint4_weight_cpu_check,
     _linear_fp_act_uint4_weight_cpu_impl,
@@ -72,6 +68,10 @@
     _linear_bf16_act_uint4_weight_check,
     _linear_bf16_act_uint4_weight_impl,
 )
+from torchao.prototype.dtypes.floatx.floatx_tensor_core_layout import (
+    _linear_f16_bf16_act_floatx_weight_check,
+    _linear_f16_bf16_act_floatx_weight_impl,
+)
 from torchao.prototype.dtypes.uintx.block_sparse_layout import (
     _linear_int8_act_int8_weight_block_sparse_check,
     _linear_int8_act_int8_weight_block_sparse_impl,