[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.

❌ 1 Awaiting Approval, 2 New Failures, 1 Pending, 4 Unrelated Failures

As of commit e6c7ba9 with merge base 81acbc1 ():

AWAITING APPROVAL - The following workflow needs approval before CI can run:

• periodic (gh)

NEW FAILURES - The following jobs have failed:

• Test WebGPU Backend / test-webgpu / test-backend-linux (webgpu, models) / linux-job (gh)
  curl: (22) The requested URL returned error:
• Test WebGPU Backend / test-webgpu / test-backend-linux (webgpu, operators) / linux-job (gh)
  curl: (22) The requested URL returned error:

BROKEN TRUNK - The following jobs failed but were present on the merge base:

👉 Rebase onto the `viable/strict` branch to avoid these failures

• pull / unittest / macos / macos-job (gh) (trunk failure)
  examples/models/muse-glimmer/tests/test_cuda_sequence_export_inputs.py::SequenceExportInputsTest::test_small_logits_example_preserves_token_and_embedding_capacity
• pull / unittest / windows / windows-job (gh) (trunk failure)
  .ci/scripts/tests/test_torchci_workflow.py::TorchciWorkflowTest::test_all_four_callers_have_the_intended_opt_in_and_permission
• trunk / unittest-release / macos / macos-job (gh) (trunk failure)
  examples/models/muse-glimmer/tests/test_cuda_sequence_export_inputs.py::SequenceExportInputsTest::test_small_logits_example_preserves_token_and_embedding_capacity
• trunk / unittest-release / windows / windows-job (gh) (trunk failure)
  .ci/scripts/tests/test_torchci_workflow.py::TorchciWorkflowTest::test_all_four_callers_have_the_intended_opt_in_and_permission

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.