Skip to content

[PyTorch] Support divisor_override for avg_pool2d and avg_pool3d #2745

Open

Description

@ash010413

Description

The PyTorch frontend currently rejects avg_pool2d and avg_pool3d when divisor_override is provided.

Minimal reproduction

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:

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:
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

Activity

  1. added
    feature requestFunctionality does not currently exist, would need to be created as a new feature (type)
    on Jun 20, 2026
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Metadata

Metadata

Assignees

No one assigned

    Labels

    feature requestFunctionality does not currently exist, would need to be created as a new feature (type)

    Type

    No type

    Projects

    No projects

      Milestone

      No milestone

      Relationships

      None yet

      Development

      No branches or pull requests

      Issue actions