Fixes

Author: jainapurva

torchao/dtypes/__init__.py

@@ -14,7 +14,6 @@
 )
 from .nf4tensor import NF4Tensor, to_nf4
 from .uintx import (
-    BlockSparseLayout,
     CutlassInt4PackedLayout,
     Int4CPULayout,
     Int4XPULayout,
@@ -29,6 +28,7 @@
     UintxLayout,
     to_marlinqqq_quantized_intx,
 )
+from .uintx.block_sparse_layout import BlockSparseLayout
 from .utils import (
     Layout,
     PlainLayout,

torchao/dtypes/affine_quantized_tensor_ops.py

@@ -25,6 +25,12 @@
     _linear_f16_bf16_act_floatx_weight_check,
     _linear_f16_bf16_act_floatx_weight_impl,
 )
+from torchao.prototype.dtypes.uintx.cutlass_int4_packed_layout import (
+    _linear_int4_act_int4_weight_cutlass_check,
+    _linear_int4_act_int4_weight_cutlass_impl,
+    _linear_int8_act_int4_weight_cutlass_check,
+    _linear_int8_act_int4_weight_cutlass_impl,
+)
 from torchao.dtypes.uintx.dyn_int8_act_int4_wei_cpu_layout import (
     _linear_int8_act_int4_weight_cpu_check,
     _linear_int8_act_int4_weight_cpu_impl,
@@ -84,13 +90,7 @@
     _linear_bf16_act_uint4_weight_check,
     _linear_bf16_act_uint4_weight_impl,
 )
-from torchao.prototype.dtypes.uintx.cutlass_int4_packed_layout import (
-    _linear_int4_act_int4_weight_cutlass_check,
-    _linear_int4_act_int4_weight_cutlass_impl,
-    _linear_int8_act_int4_weight_cutlass_check,
-    _linear_int8_act_int4_weight_cutlass_impl,
-)
-from torchao.prototype.sparsity.block_sparse_layout import (
+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,
 )

torchao/dtypes/uintx/__init__.py

@@ -1,6 +1,3 @@
-from .block_sparse_layout import (
-    BlockSparseLayout,
-)
 from .cutlass_int4_packed_layout import (
     CutlassInt4PackedLayout,
 )
@@ -39,7 +36,6 @@
 
 __all__ = [
     "UintxLayout",
-    "BlockSparseLayout",
     "MarlinSparseLayout",
     "SemiSparseLayout",
     "TensorCoreTiledLayout",