apple / apple/coremltools

[PyTorch] Support divisor_override for avg_pool2d and avg_pool3d

Open
#2,745 0 comments 0 reactions 0 assignees View on GitHub
feature request
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

Open the contributing 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

Get new issues in your inbox

A short digest of beginner-friendly GitHub issues.