Merge branch 'main' into aqt_refactor

Author: jainapurva

torchao/dtypes/__init__.py

@@ -1,20 +1,9 @@
-from .nf4tensor import NF4Tensor, to_nf4
-
-# from ..prototype.dtypes.uint2 import UInt2Tensor, BitnetTensor
-from .uintx import UInt4Tensor
 from .affine_quantized_tensor import (
     AffineQuantizedTensor,
-    to_affine_quantized_intx,
-    to_affine_quantized_intx_static,
-    # experimental, will be merged into floatx in the future
-    to_affine_quantized_fpx,
-    to_affine_quantized_floatx,
-    to_affine_quantized_floatx_static,
-)
-
-from . import affine_quantized_tensor_ops
-from .utils import (
+    Float8AQTTensorImpl,
+    Float8Layout,
     Layout,
+    MarlinSparseLayout,
     PlainLayout,
 )
 from .floatx import (
@@ -31,10 +20,17 @@
     UInt4Tensor,
     SemiSparseLayout,
     TensorCoreTiledLayout,
-    MarlinSparseLayout,
-    PlainAQTTensorImpl,
-    BlockSparseLayout,
+    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 .nf4tensor import NF4Tensor, to_nf4
+
+# from ..prototype.dtypes.uint2 import UInt2Tensor, BitnetTensor
+from .uint4 import UInt4Tensor
 
 __all__ = [
     "NF4Tensor",

torchao/dtypes/floatx/floatx_tensor_core_layout.py

@@ -1,5 +1,6 @@
+from dataclasses import dataclass
 from functools import reduce
-from typing import Tuple, Optional
+from typing import Optional, Tuple
 
 import torch
 from torch import Tensor