[EylonKrause]

Summary

DecomposeAvgPool2dPass lowers avg_pool2d with a divisor_override to avg_pool2d(count_include_pad=False) on a zero padded input, followed by a multiplication by kernel_size / divisor_override. That is only exact when every window is full. _compute_post_pad returned early for padding=0, so with ceil_mode=True the last, partial window was averaged over its valid elements only.

Example: avg_pool2d(kernel=(1, 2), stride=(1, 2), padding=0, ceil_mode=True, divisor_override=2) on [1, 2, 3, 4, 5]: PyTorch gives [1.5, 3.5, 2.5] (5 / 2), the decomposition gave [1.5, 3.5, 5.0]. This is a silent numerical error in the lowered model.

This change computes the post padding for ceil_mode with a divisor_override also when padding=0, and inserts the pad whenever any pad value is non-zero rather than only when the symmetric pre-padding is. Behaviour only changes when padding=0, ceil_mode=True and divisor_override are all set.

Independent of #24893 (it touches a different function in the same file; the two branches merge cleanly).

Release notes: arm

Test plan

• New cases divisor_override_ceil_mode_no_padding and divisor_override_ceil_mode_no_padding_count_include_pad in test_avg_pool2d.py; they fail on main in the TOSA FP and INT variants and pass with this change.
• Sweep over kernel 2-3, stride 1-3, padding 0-1, ceil_mode, count_include_pad, divisor_override in {None, 1, 2, 4, 7}, input height 4-15: 176 configurations with padding=0 fail on main (all with ceil_mode and a divisor_override), 0 with this change.
• test_avg_pool2d.py, test_decompose_avg_pool2d_pass.py, test_rewrite_avg_pool2d_pass.py, test_adaptive_avg_pool2d.py, test_mean_dim.py (excluding u55/u85/vgf): no new failures.
• Not run locally: the u55/u85/vgf variants (no Corstone FVP / Vela installed here), so those rely on CI.
• ufmt check and flake8 clean on the changed files.

cc @digantdesai @freddan80 @per @zingo @oscarandersson8218 @mansnils @Sebastian-Larsson @robell @rascani @nil-is-all

[pytorch-bot]

🔗 Helpful Links

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

• 📄 Preview Python docs built from this PR

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

✅ No Failures

As of commit 7fb7fb7 with merge base b3b565b ():
💚 Looks good so far! There are no failures yet. 💚

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

[linux-foundation-easycla]

The committers listed above are authorized under a signed CLA.

• ✅ login: EylonKrause / name: Eylon Eliyahu Krause (e6c7ba9)

[pytorch-bot]

Workflows were awaiting approval. CI has now been triggered for the ciflow labels on this PR.

[zingo]

Thanks a lot for this PRs :) I kicked of some extended testing here also. This job also ask you for EasyCLA but you onl;y need to to it once and I think the CI will pick it up automatically if not a simple rebase will usually help, lets see.

[EylonKrause]

/easycla

[EylonKrause]

@pytorchbot label "release notes: arm"

[EylonKrause]

Thanks @zingo! The CLA is signed (EasyCLA passed on #24893 for the same author), so I've asked EasyCLA to re-check this PR.

[zingo]

Thanks a lot for the fixes :) OK To merge! I'm not sure if you have access, if not let me know and I press the button :)