diff --git a/backends/arm/_passes/decompose_avg_pool2d_pass.py b/backends/arm/_passes/decompose_avg_pool2d_pass.py index e0c974598bc..f19c4fb515c 100644 --- a/backends/arm/_passes/decompose_avg_pool2d_pass.py +++ b/backends/arm/_passes/decompose_avg_pool2d_pass.py @@ -46,7 +46,12 @@ def _compute_post_pad( divisor_override, ) -> int | torch.SymInt: - if pad == 0: + # Without padding, count_include_pad=True divides a window that runs past + # the end of the input by its valid elements only, as the decomposed + # count_include_pad=False avg_pool2d does. A divisor_override needs full + # windows though, so the part of a ceil_mode window outside the input must + # still be padded with zeros. + if pad == 0 and not (ceil_mode and divisor_override is not None): return pad pad_adjust = adjust_pooling_pad_if_needed(size, kernel, stride, pad, ceil_mode) @@ -141,7 +146,7 @@ def call_operator(self, op, args, kwargs, meta): divisor_override, ) - if count_include_pad and (pad_h > 0 or pad_w > 0): + if count_include_pad and any(p != 0 for p in pad): if op in aten_avg_pool2d: pad_args = (x, pad, "constant", 0.0) else: diff --git a/backends/arm/test/ops/test_avg_pool2d.py b/backends/arm/test/ops/test_avg_pool2d.py index 9d2e863a31e..7ec9e3cb77d 100644 --- a/backends/arm/test/ops/test_avg_pool2d.py +++ b/backends/arm/test/ops/test_avg_pool2d.py @@ -110,6 +110,14 @@ def forward(self, x: torch.Tensor): AvgPool2d(3, 2, 1, True, True, divisor_override=2), (torch.rand(1, 1, 14, 14),), ), + "divisor_override_ceil_mode_no_padding": lambda: ( + AvgPool2d(3, 2, 0, True, False, divisor_override=2), + (torch.rand(1, 1, 15, 14),), + ), + "divisor_override_ceil_mode_no_padding_count_include_pad": lambda: ( + AvgPool2d(3, 2, 0, True, True, divisor_override=4), + (torch.rand(1, 1, 14, 14),), + ), "non_divisible_no_padding": lambda: ( AvgPool2d(3, 2, 0), (torch.rand(1, 16, 56, 56),),