Experiencing "divided by two must all be the same length" when trying to convert a PyTorch MaxPool2d

[anthonydito]

🐞Describing the bug

• Getting ValueError: input_shape (length 1), kernel_shape (length 2), strides (length 2), dilations (length 2), and custom_pad (length 4) divided by two must all be the same length
• In the contrived example plus my real model everything works correctly. I would except coremltools to be able to convert this layer.

Stack Trace

Error
Traceback (most recent call last):
  File "/Users/anthonydito/dev/monorepo/py/test_utils/__init__.py", line 21, in wrapper
    return func(*args, **kwargs)
  File "/Users/anthonydito/dev/monorepo/py/ml/models/shared/test_dito_model.py", line 14, in test_coreml_conversion_works
    model.save_coreml_to_s3(context)
  File "/Users/anthonydito/dev/monorepo/py/ml/models/shared/dito_model.py", line 49, in save_coreml_to_s3
    model = ct.converters.convert(
  File "/Users/anthonydito/dev/monorepo/py/.venv/lib/python3.9/site-packages/coremltools/converters/_converters_entry.py", line 574, in convert
    mlmodel = mil_convert(
  File "/Users/anthonydito/dev/monorepo/py/.venv/lib/python3.9/site-packages/coremltools/converters/mil/converter.py", line 188, in mil_convert
    return _mil_convert(model, convert_from, convert_to, ConverterRegistry, MLModel, compute_units, **kwargs)
  File "/Users/anthonydito/dev/monorepo/py/.venv/lib/python3.9/site-packages/coremltools/converters/mil/converter.py", line 212, in _mil_convert
    proto, mil_program = mil_convert_to_proto(
  File "/Users/anthonydito/dev/monorepo/py/.venv/lib/python3.9/site-packages/coremltools/converters/mil/converter.py", line 286, in mil_convert_to_proto
    prog = frontend_converter(model, **kwargs)
  File "/Users/anthonydito/dev/monorepo/py/.venv/lib/python3.9/site-packages/coremltools/converters/mil/converter.py", line 108, in __call__
    return load(*args, **kwargs)
  File "/Users/anthonydito/dev/monorepo/py/.venv/lib/python3.9/site-packages/coremltools/converters/mil/frontend/torch/load.py", line 80, in load
    return _perform_torch_convert(converter, debug)
  File "/Users/anthonydito/dev/monorepo/py/.venv/lib/python3.9/site-packages/coremltools/converters/mil/frontend/torch/load.py", line 99, in _perform_torch_convert
    prog = converter.convert()
  File "/Users/anthonydito/dev/monorepo/py/.venv/lib/python3.9/site-packages/coremltools/converters/mil/frontend/torch/converter.py", line 519, in convert
    convert_nodes(self.context, self.graph)
  File "/Users/anthonydito/dev/monorepo/py/.venv/lib/python3.9/site-packages/coremltools/converters/mil/frontend/torch/ops.py", line 88, in convert_nodes
    add_op(context, node)
  File "/Users/anthonydito/dev/monorepo/py/.venv/lib/python3.9/site-packages/coremltools/converters/mil/frontend/torch/ops.py", line 1405, in max_pool2d
    _max_pool(context, node, inputs)
  File "/Users/anthonydito/dev/monorepo/py/.venv/lib/python3.9/site-packages/coremltools/converters/mil/frontend/torch/ops.py", line 1379, in _max_pool
    pool = mb.max_pool(
  File "/Users/anthonydito/dev/monorepo/py/.venv/lib/python3.9/site-packages/coremltools/converters/mil/mil/ops/registry.py", line 182, in add_op
    return cls._add_op(op_cls_to_add, **kwargs)
  File "/Users/anthonydito/dev/monorepo/py/.venv/lib/python3.9/site-packages/coremltools/converters/mil/mil/builder.py", line 184, in _add_op
    new_op.type_value_inference()
  File "/Users/anthonydito/dev/monorepo/py/.venv/lib/python3.9/site-packages/coremltools/converters/mil/mil/operation.py", line 257, in type_value_inference
    output_types = self.type_inference()
  File "/Users/anthonydito/dev/monorepo/py/.venv/lib/python3.9/site-packages/coremltools/converters/mil/mil/ops/defs/iOS15/pool.py", line 69, in type_inference
    D_out_shape = spatial_dimensions_out_shape(
  File "/Users/anthonydito/dev/monorepo/py/.venv/lib/python3.9/site-packages/coremltools/converters/mil/mil/ops/defs/_utils.py", line 285, in spatial_dimensions_out_shape
    raise ValueError(
ValueError: input_shape (length 1), kernel_shape (length 2), strides (length 2), dilations (length 2), and custom_pad (length 4) divided by two must all be the same length

/Applications/Xcode.app/Contents/Developer/Library/Frameworks/Python3.framework/Versions/3.9/lib/python3.9/tempfile.py:817: ResourceWarning: Implicitly cleaning up <TemporaryDirectory '/var/folders/cx/bvq8nwgx05503mzx92jfnh0c0000gn/T/tmp3zbmhu96'>
  _warnings.warn(warn_message, ResourceWarning)


Ran 1 test in 0.247s

FAILED (errors=1)

Process finished with exit code 1

To Reproduce

Reproduce the bug: https://colab.research.google.com/drive/1azCE_B8Eu9D9GV72W4fSwRLJlZuf2ij3?usp=sharing

import torch
import coremltools as ct

class CoreMLTestModel(torch.nn.Module):

    def __init__(self):
        super(CoreMLTestModel, self).__init__()
        self.max_pool1 = torch.nn.MaxPool2d(kernel_size=3, stride=1, padding=1)

    def forward(self, x):
        x = self.max_pool1(x)
        return x

model = CoreMLTestModel()
model.eval()
input_size = torch.Size([2, 19, 30])
example_input = torch.rand(input_size)
print(model(example_input))
traced_model = torch.jit.trace(model, example_input)

# error here
result = ct.convert(
    traced_model,
    minimum_deployment_target=ct.target.iOS17,
    inputs=[ct.TensorType(shape=example_input.shape)]
)

System environment (please complete the following information):

• coremltools version: 7.1
• OS (e.g. MacOS version or Linux type): MacOS 14.2 Beta and whatever Google Collab uses
• Any other relevant version information (e.g. PyTorch or TensorFlow version): Tested on PyTorch 2.1.0+cu121 and PyTorch 2.3.0.dev20240212

Additional context

• Nothing additional. Please let me know if you have any questions