up

Author: metascroy

backends/xnnpack/partition/config/quant_affine_configs.py

@@ -36,7 +36,7 @@ def get_original_aten(self) -> Optional[torch._ops.OpOverload]:
         try:
             import torchao.quantization.quant_primitives  # noqa
 
-            return torch.ops.quant.quantize_affine.default
+            return torch.ops.torchao.quantize_affine.default
         except:
             return None
 
@@ -48,7 +48,7 @@ def get_original_aten(self) -> Optional[torch._ops.OpOverload]:
         try:
             import torchao.quantization.quant_primitives  # noqa
 
-            return torch.ops.quant.dequantize_affine.default
+            return torch.ops.torchao.dequantize_affine.default
         except:
             return None
 
@@ -60,6 +60,6 @@ def get_original_aten(self) -> Optional[torch._ops.OpOverload]:
         try:
             import torchao.quantization.quant_primitives  # noqa
 
-            return torch.ops.quant.choose_qparams_affine.default
+            return torch.ops.torchao.choose_qparams_affine.default
         except:
             return None

backends/xnnpack/test/ops/test_linear.py

@@ -402,9 +402,9 @@ def _test_groupwise_dq_linear(
             .export()
             .check_count(
                 {
-                    "torch.ops.quant.choose_qparams_affine.default": 1 * num_linears,
-                    "torch.ops.quant.quantize_affine.default": 1 * num_linears,
-                    "torch.ops.quant.dequantize_affine.default": 2 * num_linears,
+                    "torch.ops.torchao.choose_qparams_affine.default": 1 * num_linears,
+                    "torch.ops.torchao.quantize_affine.default": 1 * num_linears,
+                    "torch.ops.torchao.dequantize_affine.default": 2 * num_linears,
                     "torch.ops.aten.linear.default": 1 * num_linears,
                 }
             )