Decompose aten.var_mean.correction in to_edge

[usmanmateen]

Summary

Fixes #23680.

torch.var_mean and torch.std_mean export to aten.var_mean.correction. It isn't a core ATen op and isn't in core_aten_decompositions(), so to_edge rejected it with SpecViolationError on the portable and XNNPACK paths. Its only registered decomposition is the torch._refs one, which produces prims ops. That's why this adds a small decomposition to _default_decomposition_table instead of listing the op in additional_decomp_ops (that gives Missing out variants: {'prims::var', 'prims::sum', 'prims::div'}). The decomposition goes to aten.var.correction and aten.mean.dim, both core ATen ops with portable kernels.

Test plan

test_to_edge_decomposes_var_mean_and_std_mean in exir/program/test/test_program.py exports a module that calls both ops, runs it through the portable pybindings, and compares with eager. It fails with the SpecViolationError before this change and passes after.

I also lowered and ran these cases with to_edge and with to_edge_transform_and_lower(..., partitioner=[XnnpackPartitioner()]), comparing outputs with eager (atol=1e-5):

case	before	after (portable / XNNPACK)
torch.var_mean(x, dim=1, correction=0)	SpecViolationError	match / match
torch.var_mean(x) (dim=None)	SpecViolationError	match / match
torch.var_mean(x, dim=(0, 2), keepdim=True)	SpecViolationError	match / match
torch.var_mean(x, dim=-1, correction=2)	SpecViolationError	match / match
torch.std_mean(x, dim=1, unbiased=False)	SpecViolationError	match / match

Tested against executorch 1.5.1 and torch 2.14.1+cpu on Linux x86_64 (WSL2); _default_decomposition_table is the same on main. A set of 55 other small modules (activations, reductions, norms, convolutions, pooling, RNNs, attention, and six torchvision models) lowers and matches eager on both paths before and after. ufmt (black 24.4.2, usort 1.0.8.post1) and flake8 pass on the changed files.

This PR was authored with AI assistance (Claude Code).

cc @nil-is-all

[pytorch-bot]

🔗 Helpful Links

🧪 See artifacts and rendered test results at hud.pytorch.org/pr/pytorch/executorch/23681

• 📄 Preview Python docs built from this PR

Note: Links to docs will display an error until the docs builds have been completed.

⚠️ 15 Awaiting Approval

As of commit 9e89edd with merge base f2575be ():

AWAITING APPROVAL - The following workflows need approval before CI can run:

• Build documentation (gh)
• Build Presets (gh)
• Cadence Build & Test (gh)
• Lint (gh)
• periodic (gh)
• pull (gh)
• Test ARM Backend (gh)
• Test CoreML Backend (gh)
• Test Cortex-M Backend (gh)
• Test NXP Backend (gh)
• Test QNN Backend (gh)
• Test QNN Windows MSVC build (gh)
• Test Vulkan Backend (gh)
• Test WebGPU Backend (gh)
• Test XNNPACK Backend (gh)

This comment was automatically generated by Dr. CI and updates every 15 minutes.

[meta-cla]

Hi @usmanmateen!

Thank you for your pull request and welcome to our community.

Action Required

In order to merge any pull request (code, docs, etc.), we require contributors to sign our Contributor License Agreement, and we don't seem to have one on file for you.

Process

In order for us to review and merge your suggested changes, please sign at https://code.facebook.com/cla. If you are contributing on behalf of someone else (eg your employer), the individual CLA may not be sufficient and your employer may need to sign the corporate CLA.

Once the CLA is signed, our tooling will perform checks and validations. Afterwards, the pull request will be tagged with CLA signed. The tagging process may take up to 1 hour after signing. Please give it that time before contacting us about it.

If you have received this in error or have any questions, please contact us at cla@meta.com. Thanks!

[linux-foundation-easycla]

The committers listed above are authorized under a signed CLA.

• ✅ login: usmanmateen / name: usmanmateen (9e89edd)

[usmanmateen]

@pytorchbot label "release notes: exir"

[meta-cla]

Thank you for signing our Contributor License Agreement. We can now accept your code for this (and any) Meta Open Source project. Thanks!