Refactor Affine Quantized Tensor

test/dtypes/test_affine_quantized.py

@@ -92,7 +92,7 @@ def test_to_device(self, apply_quant):
 
     @unittest.skipIf(not torch.cuda.is_available(), "Need CUDA available")
     def test_register_new_dispatch(self):
-        from torchao.dtypes.affine_quantized_tensor import (
+        from torchao.dtypes.affine_quantized_tensor_ops import (
             register_aqt_quantized_linear_dispatch,
             deregister_aqt_quantized_linear_dispatch,
         )

test/dtypes/test_floatx.py

@@ -15,7 +15,7 @@
     to_scaled_tc_floatx,
     from_scaled_tc_floatx,
 )
-from torchao.dtypes.floatx.floatx import _pack_tc_floatx, _pack_tc_fp6
+from torchao.dtypes.floatx.floatx_tensor_core_layout import _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 (
     quantize_,

test/dtypes/test_uint4.py

@@ -1,5 +1,5 @@
 import torch
-from torchao.dtypes.uint4 import (
+from torchao.dtypes.uintx.uint4_layout import (
     UInt4Tensor,
     PerChannelSymmetricWeightUInt4Tensor,
 )

torchao/dtypes/__init__.py

@@ -1,23 +1,40 @@
+from . import affine_quantized_tensor_ops
+
+# from ..prototype.dtypes.uint2 import UInt2Tensor, BitnetTensor
 from .affine_quantized_tensor import (
     AffineQuantizedTensor,
-    Float8AQTTensorImpl,
-    Float8Layout,
-    Layout,
-    MarlinSparseLayout,
-    PlainLayout,
-    SemiSparseLayout,
-    TensorCoreTiledLayout,
     to_affine_quantized_floatx,
     to_affine_quantized_floatx_static,
     # experimental, will be merged into floatx in the future
     to_affine_quantized_fpx,
     to_affine_quantized_intx,
     to_affine_quantized_intx_static,
 )
+from .floatx import (
+    Float8AQTTensorImpl,
+    Float8Layout,
+)
 from .nf4tensor import NF4Tensor, to_nf4
+from .uintx import (
+    _BIT_WIDTH_TO_DTYPE,
+    _DTYPE_TO_BIT_WIDTH,
+    BlockSparseLayout,
+    MarlinSparseLayout,
+    PlainAQTTensorImpl,
+    SemiSparseLayout,
+    TensorCoreTiledLayout,
+    UInt4Tensor,
+    UintxAQTTensorImpl,
+    UintxLayout,
+    UintxTensor,
+    to_uintx,
+)
+from .utils import (
+    Layout,
+    PlainLayout,
+)
 
 # from ..prototype.dtypes.uint2 import UInt2Tensor, BitnetTensor
-from .uint4 import UInt4Tensor
 
 __all__ = [
     "NF4Tensor",
@@ -36,4 +53,15 @@
     "Float8Layout",
     "Float8AQTTensorImpl",
     "MarlinSparseLayout",
+    "PlainAQTTensorImpl",
+    "affine_quantized_tensor_ops",
+    "BlockSparseLayout",
+    "to_uintx",
+    "UintxTensor",
+    "UintxLayout",
+    "UintxAQTTensorImpl",
+    "_DTYPE_TO_BIT_WIDTH",
+    "_BIT_WIDTH_TO_DTYPE",
+    "Uint4Tensor",
+    "PlainAQTTensorImpl",
 ]