Skip to content

Fix avg_pool count_include_pad argument guard - #80

Open
devin-lai wants to merge 1 commit into
apple:mainfrom
devin-lai:fix/avgpool-count-include-pad
Open

Fix avg_pool count_include_pad argument guard#80
devin-lai wants to merge 1 commit into
apple:mainfrom
devin-lai:fix/avgpool-count-include-pad

Conversation

@devin-lai

Copy link
Copy Markdown

Fixes #73.

torch.export can produce five-argument average-pool nodes when ceil_mode is set but count_include_pad remains at its default. The lowering checked argument 4 before reading argument 5, which caused an IndexError.

This updates both the 2D and 3D lowerings to guard argument 5 directly and adds end-to-end coverage for the omitted default.

Tested with pytest -q tests/ops/test_ops.py::test_avg_pool_defaults_count_include_pad_when_omitted.

Fixes apple#73.

ATen omits trailing default arguments from exported nodes. When ceil_mode is true and count_include_pad remains at its default, the node contains exactly five arguments. The lowering guarded args[5] with a check for args[4], which caused an IndexError.

Guard index 5 directly for both avg_pool2d and avg_pool3d while retaining the True default, and add end-to-end regression coverage for both operators.
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

None yet

Projects

None yet

Development

Successfully merging this pull request may close these issues.

avg_pool2d: count_include_pad reads node.args[5] but guards node.args[4] — IndexError on 5-arg nodes (ceil_mode=True), still present in 0.4.2

1 participant