Missing out variants: {'torchao::dequantize_affine'}

[guangy10]

🐛 Describe the bug

Hitting this issue when enabling quant embedding using torchao API on native Hugging Face models in optimum-executorch.

RuntimeError: Node torchao_dequantize_affine_default with op <EdgeOpOverload: torchao.dequantize_affine.default>: schema = torchao::dequantize_affine(Tensor input, SymInt[] block_size, Tensor scale, Tensor? zero_point, ScalarType input_dtype, Scalar? quant_min=None, Scalar? quant_max=None, str? zero_point_domain="INT", ScalarType output_dtype=6) -> Tensor was not decomposed or delegated.
This op was registered by the partitioner XnnpackPartitioner to not be decomposed.
The following ops: {<OpOverload(op='aten.linear', overload='default')>, <OpOverload(op='aten.prelu', overload='default')>, <OpOverload(op='torchao.quantize_affine', overload='default')>, <OpOverload(op='aten.max_pool2d', overload='default')>, <OpOverload(op='torchao.dequantize_affine', overload='default')>, <OpOverload(op='aten.upsample_bilinear2d', overload='vec')>, <OpOverload(op='torchao.choose_qparams_affine', overload='default')>} were specified to not be decomposed in forward.

Find the stracktrace in the CI: https://github.com/huggingface/optimum-executorch/actions/runs/14801963348/job/41562587218?pr=62#step:5:1502

Versions

Versions of relevant libraries:
[pip3] executorch==0.7.0a0+3e38b85
[pip3] flake8==6.1.0
[pip3] flake8-breakpoint==1.1.0
[pip3] flake8-bugbear==24.4.26
[pip3] flake8-comprehensions==3.14.0
[pip3] flake8-plugin-utils==1.3.3
[pip3] flake8-pyi==23.5.0
[pip3] mypy==1.14.1
[pip3] mypy-extensions==1.0.0
[pip3] numpy==2.2.4
[pip3] optimum-executorch==0.0.0.dev0
[pip3] pytorch_tokenizers==0.1.0
[pip3] torch==2.8.0.dev20250325
[pip3] torchao==0.11.0.dev20250422+cpu
[pip3] torchaudio==2.6.0.dev20250325
[pip3] torchsr==1.0.4
[pip3] torchvision==0.22.0.dev20250325
[conda] executorch 0.7.0a0+3e38b85 pypi_0 pypi
[conda] numpy 2.2.4 pypi_0 pypi
[conda] optimum-executorch 0.0.0.dev0 pypi_0 pypi
[conda] pytorch-tokenizers 0.1.0 pypi_0 pypi
[conda] torch 2.8.0.dev20250325 pypi_0 pypi
[conda] torchao 0.11.0.dev20250422+cpu pypi_0 pypi
[conda] torchaudio 2.6.0.dev20250325 pypi_0 pypi
[conda] torchfix 0.6.0 pypi_0 pypi
[conda] torchsr 1.0.4 pypi_0 pypi
[conda] torchvision 0.22.0.dev20250325 pypi_0 pypi

cc @digantdesai @mcr229 @cbilgin @larryliu0820 @manuelcandales

export_prog.txt

[guangy10]

cc: @jerryzh168 @metascroy @mcr229

[metascroy]

torchao::dequantize_affine is not suppose to have an out variant and it should not be left in the graph after quantization and passes run.

Can you share a print out of the exported_program after lowering to XNNPACK, but before to_executorch?

[guangy10]

@metascroy export graph is attached

[metascroy]

The issue is the embedding has padding_idx, which is not recognized by quant_patterns and we do not have a quantized kernel for it. But after training, padding_idx doesn't matter much, so you could write a pass to just strip it out, e.g., RemoveEmbeddingPaddingIdxPass, which does:

torch.ops.aten.embedding(idx, weight, padding_idx=x) -> torch.ops.embedding(idx, weight)