avg_pool1d/avg_pool2d with ceil_mode=True silently return incorrect values

[LeSingh1]

🐞 Describing the bug

avg_pool1d / avg_pool2d with ceil_mode=True silently produce wrong numbers. The model converts and predicts without any error or warning.

_avg_pool forwards count_include_pad straight into MIL's exclude_padding_from_average, but the two do not mean the same thing for ceil mode: MIL counts the ceil-mode overflow region as padding, and PyTorch never does. With PyTorch's default count_include_pad=True, every ceil-extended window ends up divided by the full kernel size instead of the number of real elements in it.

avg_pool3d already raises for this combination (ops.py, the ceil_mode + count_include_pad guard), so the 3D path treats it as unsupported while 1D and 2D return silently incorrect results.

Stack Trace

None — there is no error. That is the problem.

To Reproduce

import numpy as np, torch, torch.nn as nn, coremltools as ct

class M(nn.Module):
    def forward(self, x):
        return torch.nn.functional.avg_pool2d(x, 2, stride=2, ceil_mode=True)

x = torch.arange(9, dtype=torch.float32).reshape(1, 1, 3, 3)
tm = torch.jit.trace(M().eval(), x)
mlm = ct.convert(
    tm,
    inputs=[ct.TensorType(shape=x.shape)],
    convert_to="mlprogram",
    compute_units=ct.ComputeUnit.CPU_ONLY,
    minimum_deployment_target=ct.target.iOS16,
)
out = list(mlm.predict({list(mlm.input_description)[0]: x.numpy()}).values())[0]

print("torch  :", M()(x).flatten().tolist())
print("coreml :", np.array(out).flatten().tolist())

torch  : [2.0, 3.5, 6.5, 8.0]
coreml : [2.0, 1.75, 3.25, 2.0]

Three of the four outputs are wrong. Only the top-left window — the one that needs no ceil extension — is correct. The others are the right sums divided by 4 instead of by 2, 2 and 1.

System environment

• coremltools version: main @ a5a17b2
• OS: macOS (Apple silicon)
• torch: 2.8.0
• Python: 3.13

Additional context

The existing TestAvgPool coverage only ever pairs ceil_mode=True with count_include_pad=True when kernel_size=1, where no window can be ceil-extended — so this combination is effectively untested.

I have a fix working locally, but the shape of it is a design call I would rather not make unilaterally: the natural approach is to hoist explicit padding into a separate mb.pad so that real padding still counts toward the average while the ceil overflow does not. That touches the same area as #2748, so I did not want to send a competing diff without a maintainer's view.

Happy to open a PR — just let me know whether you would prefer that direction, or the avg_pool3d treatment of raising for the unsupported combination, and I will follow it.

[TobyRoseman]

I can't reproduce this. The output matches for me:

torch  : [2.0, 3.5, 6.5, 8.0]
coreml : [2.0, 3.5, 6.5, 8.0]

I'm using the same commit, same version of PyTorch and same version of Python.