Rebase and lint fixes

Author: jainapurva

test/hqq/test_hqq_affine.py

@@ -1,12 +1,7 @@
 import unittest
 import torch
 from torchao.dtypes.affine_quantized_tensor import (
-    to_affine_quantized_intx,
     ZeroPointDomain,
-    PlainAQTTensorImpl,
-    PlainLayout,
-    TensorCoreTiledAQTTensorImpl,
-    TensorCoreTiledLayout,
     MappingType,
 )
 

torchao/dtypes/affine_quantized_tensor.py

@@ -140,6 +140,8 @@ def dequantize(self, output_dtype: Optional[torch.dtype] = None) -> torch.Tensor
                 self.zero_point_domain,
                 output_dtype=output_dtype,
             )
+            from torchao.dtypes.uintx import TensorCoreTiledLayout
+
             if isinstance(self._layout, TensorCoreTiledLayout):
                 # need to return to original shape if tensor was padded
                 # in preprocessing

torchao/dtypes/affine_quantized_tensor_ops.py

@@ -101,11 +101,11 @@ def _quantized_linear_op(input_tensor, weight_tensor, bias):
 AffineQuantizedTensor._quantized_linear_op = _quantized_linear_op
 
 
-# # following are a list of (dispatch_condition, implementation) functions that takes the following args:
-# # input_tensor: dimension is (M1, M2, ..., in_features)
-# # weight_tensor: dimension is (out_features, in_features)
-# # bias: dimension is (out_features,)
-# # so that these can be shared by F.linear, aten.mm, aten.addmm dispatches
+# _register_aqt_quantized_linear_dispatches function has a list of (dispatch_condition, implementation) functions, defined in their dtype layout classes, that takes the following args:
+# input_tensor: dimension is (M1, M2, ..., in_features)
+# weight_tensor: dimension is (out_features, in_features)
+# bias: dimension is (out_features,)
+# so that these can be shared by F.linear, aten.mm, aten.addmm dispatches
 def _register_aqt_quantized_linear_dispatches():
     for dispatch_condition, impl in [
         (_linear_int8_act_int8_weight_check, _linear_int8_act_int8_weight_impl),

torchao/dtypes/uintx/plain_layout.py

@@ -11,9 +11,11 @@
     register_layout,
 )
 from torchao.dtypes.utils import AQTTensorImpl, Layout, PlainLayout
+from torchao.kernel import (
+    int_scaled_matmul,
+)
 from torchao.quantization.quant_primitives import (
     ZeroPointDomain,
-    int_scaled_matmul,
 )
 from torchao.utils import fill_defaults