Lint fixes;

Author: jainapurva

ruff.toml

@@ -2,6 +2,8 @@
 # We plan to add files in chunks using the 'include' list below.
 # To add a new path: Simply add it to the 'include' list.
 # Example: To lint all files in every subfolder of 'test', add "test/**/*"
+# To exclude a file type: Simply add it to the 'include' list.
+# Example: To lint all files in every subfolder of 'test', add "test/**/*"
 include = [
     "torchao/float8/inference.py",
     "torchao/float8/float8_utils.py",
@@ -10,4 +12,9 @@ include = [
     "torchao/float8/float8_tensor.py",
     "torchao/quantization/linear_activation_weight_observer.py",
     "test/quantization/test_observer.py",
+    "torchao/dtypes/*"
+]
+
+exclude = [
+    "**/*.md"
 ]

torchao/dtypes/__init__.py

@@ -21,7 +21,7 @@
 __all__ = [
     "NF4Tensor",
     "to_nf4",
-    "UInt4Tensor"
+    "UInt4Tensor",
     "AffineQuantizedTensor",
     "to_affine_quantized_intx",
     "to_affine_quantized_intx_static",

torchao/dtypes/affine_quantized_tensor.py

@@ -1,8 +1,6 @@
 import torch
 from typing import Tuple, Optional, Union
 import torchao.ops
-from collections import defaultdict
-import functools
 import math
 from torchao.quantization.quant_primitives import (
     choose_qparams_affine,
@@ -36,10 +34,10 @@
     _is_float8_type
 )
 import logging
+from torchao.float8.inference import Float8MMConfig
 
 logger = logging.getLogger(__name__)
 
-from torchao.float8.inference import Float8MMConfig
 aten = torch.ops.aten
 
 
@@ -1024,7 +1022,6 @@ def from_plain(
         packed_weight = torch.ops.aten._convert_weight_to_int4pack(int_data, layout_type.inner_k_tiles)
         scale = scale.reshape(int_data.shape[0], -1)
         zero_point = zero_point.reshape(int_data.shape[0], -1)
-        from torchao.quantization.utils import pack_tinygemm_scales_and_zeros
         scale_and_zero = pack_tinygemm_scales_and_zeros(scale, zero_point)
         return cls(packed_weight, scale_and_zero, False, layout_type)
 
@@ -1232,7 +1229,7 @@ def _linear_bf16_act_uint4_weight_check(input_tensor, weight_tensor, bias):
 
 
 def _linear_bf16_act_uint4_weight_impl(input_tensor, weight_tensor, bias):
-    assert weight_tensor.block_size[0] == 1, f"Requires groupwise quantization, got block_size: {block_size}"
+    assert weight_tensor.block_size[0] == 1, f"Requires groupwise quantization, got block_size: {weight_tensor.block_size}"
     assert input_tensor.shape[-1] == weight_tensor.shape[1], (
         f"need input_tensor shape: {input_tensor.shape} final"
         f"dim to match weight_tensor shape: {weight_tensor.shape} second dim "
@@ -1292,7 +1289,6 @@ def _linear_fp_act_int8_weight_impl(input_tensor, weight_tensor, bias):
     # per channel int8 weight only quantizated mm
     w_vals_int8_t = weight_tensor.layout_tensor.int_data.t()
     scale = weight_tensor.layout_tensor.scale
-    orig_dtype = input_tensor.dtype
     m = torch.mm(
         input_tensor.reshape(-1, input_tensor.shape[-1]),
         w_vals_int8_t.to(input_tensor.dtype),
@@ -1424,7 +1420,7 @@ def _linear_fp_act_int4_weight_sparse_marlin_check(input_tensor, weight_tensor,
     )
 
 def _linear_fp_act_int4_weight_sparse_marlin_impl(input_tensor, weight_tensor, bias):
-    from torchao.sparsity.marlin import marlin_24_workspace, const
+    from torchao.sparsity.marlin import marlin_24_workspace
     assert isinstance(weight_tensor, AffineQuantizedTensor)
 
     sparse_w_int4 = weight_tensor.layout_tensor.int_data

torchao/dtypes/fpx/__init__.py

@@ -1 +1,9 @@
 from .fpx import FpxTensorCoreLayoutType, FpxTensorCoreAQTLayout, to_scaled_tc_fpx, from_scaled_tc_fpx, _SPLIT_K_MAP
+
+__all__ = [
+    "FpxTensorCoreAQTLayout",
+    "FpxTensorCoreLayoutType",
+    "to_scaled_tc_fpx",
+    "from_scaled_tc_fpx",
+    "_SPLIT_K_MAP",
+]

torchao/dtypes/fpx/fpx.py

@@ -8,7 +8,6 @@
 from torchao.dtypes.utils import (
     LayoutType,
 )
-from torchao.quantization.quant_api import _get_linear_subclass_inserter
 from dataclasses import dataclass
 from torchao.dtypes.affine_quantized_tensor import AQTLayout, register_layout_cls
 

torchao/dtypes/uintx/bitpacking.py

@@ -160,7 +160,7 @@ def unpack_cpu(data: List[torch.Tensor],
             output_narrow = output.narrow(dim, j * group_size, group_size)
             group = data[i] & unpack_mask[bit_size][j]
             shift_amt = j * bit_size - rel_pos
-            output_narrow.copy_(torch.bitwise_or(output_narrow, abs_rsh(group, j * bit_size - rel_pos)))
+            output_narrow.copy_(torch.bitwise_or(output_narrow, abs_rsh(group, shift_amt)))
     return output
 
 # these are faster on the GPU
@@ -193,7 +193,7 @@ def _unpack(data, element_size, scale, dim):
 
     for i in range(scale):
         shift_amt = element_size * i
-        chunk = unpacked_data.narrow(dim, unpacked_data.shape[dim]*i//scale, unpacked_data.shape[dim] // scale).copy_((data >> shift_amt) & nbits)
+        unpacked_data.narrow(dim, unpacked_data.shape[dim]*i//scale, unpacked_data.shape[dim] // scale).copy_((data >> shift_amt) & nbits)
 
     return unpacked_data