[PyTorch] Support divisor_override for avg_pool2d and avg_pool3d
- Dominant language
- Python
- Stars
- 5.4k
- Forks
- 850
- Avg merge
- 4d 5h
- Merged PRs (30d)
- 10
Description
### Description
The PyTorch frontend currently rejects `avg_pool2d` and `avg_pool3d` when `divisor_override` is provided.
### Minimal reproduction
```python
import torch
import coremltools as ct
class Model(torch.nn.Module):
def forward(self, x):
return torch.nn.functional.avg_pool2d(
x,
kernel_size=2,
divisor_override=3,
)
model = Model().eval()
example = torch.rand(1, 1, 4, 4)
traced = torch.jit.trace(model, example)
ct.convert(
traced,
inputs=[ct.TensorType(shape=example.shape)],
convert_to="mlprogram",
)
```
### Current behavior
Conversion fails with:
```text
ValueError: divisor_override is not supported for avg_pool2d
```
### Expected behavior
Models using a static, nonzero `divisor_override` should convert successfully when the operation can be represented using MIL operations, while preserving numerical parity with PyTorch.
### Possible implementation
For configurations where average pooling uses a fixed full-kernel divisor:
1. Lower the operation using the existing average-pooling implementation.
2. Multiply the result by:
```text
number_of_kernel_elements / divisor_override
```
An initial implementation could support:
- `avg_pool2d` and `avg_pool3d`
- Static kernel dimensions
- Positive constant `divisor_override`
- Zero padding, or `count_include_pad=True`
Cases combining nonzero padding with `count_include_pad=False` may require separate handling because the effective divisor can vary near boundaries.
I would be happy to implement this and add regression tests if this scope is acceptable.
### Environment
- PyTorch: 2.8.0
- coremltools: 9.0
- Platform: macOS / Apple silicon
Contributor guide
Research direction
Start with the PyTorch frontend handling for avg_pool2d and avg_pool3d and reproduce the failure using the model shown in the issue. Add regression coverage for static positive divisor_override with the listed padding configurations, and verify conversion succeeds with numerical parity against PyTorch; leave boundary cases with nonzero padding and count_include_pad=False explicitly outside the initial scope.
Written by the indexing model from the issue text.
Assessment
- Tech stack
- python, pytorch
- Domain
- machine-learning, tooling
- Issue type
- Feature
- Difficulty
- 4/5
- Estimated time
- 3-5 days
- Activity status
- Quiet
- Clarity
- Mostly clear
- Newbie friendliness
- 58/100