up

Author: metascroy

test/quantization/quantize_/workflows/intx/test_intx_unpacked_tensor.py

@@ -26,12 +26,10 @@
 from torchao.quantization.granularity import PerGroup
 from torchao.quantization.quantize_.common import PackingFormat
 from torchao.quantization.utils import compute_error
-from torchao.utils import (
-    TORCH_VERSION_AT_LEAST_2_8,
-)
+from torchao.utils import torch_version_at_least
 
 
-@unittest.skipIf(not TORCH_VERSION_AT_LEAST_2_8, "Need pytorch 2.8+")
+@unittest.skipIf(not torch_version_at_least("2.8.0"), "Need pytorch 2.8+")
 class TestIntxUnpackedTensor(TestCase):
     def setUp(self):
         self.config = IntxWeightOnlyConfig(

torchao/quantization/quant_api.py

@@ -118,6 +118,7 @@
     _DTYPE_TO_QVALUE_BOUNDS,
     MappingType,
     ZeroPointDomain,
+    quantize_affine,
 )
 from .subclass import (
     QuantizedLinearWeightBase,
@@ -797,7 +798,6 @@ def _int8_dynamic_activation_intx_weight_quantize_tensor(weight, bias, config):
     act_mapping_type = config.act_mapping_type
     layout = config.layout
     packing_format = config.packing_format
-    compute_target = config.compute_target
 
     assert weight.dim() == 2, (
         f"Int8DynamicActivationIntxWeightConfig only works for 2-d Tensor, got: {weight.dim()}"

torchao/quantization/quantize_/workflows/intx/intx_unpacked_tensor.py

@@ -19,7 +19,6 @@
 )
 from torchao.quantization.utils import _get_per_token_block_size
 from torchao.utils import (
-    TORCH_VERSION_AT_LEAST_2_5,
     TorchAOBaseTensor,
     fill_defaults,
 )
@@ -300,6 +299,5 @@ def _(func, types, args, kwargs):
 
 IntxUnpackedTensor.__module__ = "torchao.quantization"
 
-if TORCH_VERSION_AT_LEAST_2_5:
-    # Allow a model with IntxUnpackedTensor weights to be loaded with `weights_only=True`
-    torch.serialization.add_safe_globals([IntxUnpackedTensor])
+# Allow a model with IntxUnpackedTensor weights to be loaded with `weights_only=True`
+torch.serialization.add_safe_globals([IntxUnpackedTensor])