From a1335161ceb2514e52597d5ffa83ce3c27749079 Mon Sep 17 00:00:00 2001 From: Emmanuel Dumont Date: Mon, 9 Mar 2026 06:32:46 -0700 Subject: [PATCH] Fix AttributeError when torchao ops are not registered Summary: The `_QUANT_PRIMITIVES` list initialization in `util.py` catches `ImportError` when torchao is not installed, but fails to catch `AttributeError` when torchao is installed but the operators (`dequantize_affine`, `quantize_affine`, `choose_qparams_affine`) are not registered. This causes import failures in environments where torchao is available but an older version is used that doesn't have these operators registered. The fix adds `AttributeError` to the exception handler alongside `ImportError`. Differential Revision: D95735505 --- exir/operator/util.py | 8 ++++++-- 1 file changed, 6 insertions(+), 2 deletions(-) diff --git a/exir/operator/util.py b/exir/operator/util.py index 23dc3edd302..fd900e6f635 100644 --- a/exir/operator/util.py +++ b/exir/operator/util.py @@ -55,7 +55,9 @@ def gen_out_variant_schema(func_op_schema: str) -> str: torch.ops.quantized_decomposed.choose_qparams.tensor, ] try: - import torchao # noqa: F401 + # Import quant_primitives directly to ensure custom ops are registered + # before accessing them via torch.ops.torchao + import torchao.quantization.quant_primitives # noqa: F401 _QUANT_PRIMITIVES.extend( [ @@ -64,5 +66,7 @@ def gen_out_variant_schema(func_op_schema: str) -> str: torch.ops.torchao.choose_qparams_affine.default, ] ) -except ImportError: +except (ImportError, AttributeError): + # ImportError: torchao or quant_primitives not installed + # AttributeError: torchao installed but operators not registered (e.g., older version) pass