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,
         )

torchao/dtypes/__init__.py

@@ -1,6 +1,6 @@
 from .nf4tensor import NF4Tensor, to_nf4
 # from ..prototype.dtypes.uint2 import UInt2Tensor, BitnetTensor
-from .uint4 import UInt4Tensor
+from .uintx import UInt4Tensor
 from .affine_quantized_tensor import (
     AffineQuantizedTensor,
     to_affine_quantized_intx,
@@ -9,15 +9,22 @@
     to_affine_quantized_fpx,
     to_affine_quantized_floatx,
     to_affine_quantized_floatx_static,
+    PlainAQTTensorImpl,
+)
+from .affine_quantized_tensor_ops import *
+from .utils import (
     Layout,
     PlainLayout,
-    SemiSparseLayout,
-    TensorCoreTiledLayout,
+)
+from .floatx import (
     Float8Layout,
     Float8AQTTensorImpl,
+)
+from .uintx import (
+    SemiSparseLayout,
+    TensorCoreTiledLayout,
     MarlinSparseLayout,
 )
-
 __all__ = [
     "NF4Tensor",
     "to_nf4",