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
- Dominant language
- Python
- Stars
- 152
- Forks
- 45
- Avg merge
- 1d 7m
- Merged PRs (30d)
- 12
Description
### Summary
`replace_avg_pool2d` in `coreai_torch/_aten_to_core.py` reads `count_include_pad` from
`node.args[5]` but guards on `node.args[4]`. A node carrying **exactly five** arguments passes
the guard and then raises `IndexError: tuple index out of range`, so conversion fails before any
tensor is produced.
Reproduced on **coreai-torch 0.4.1 and 0.4.2** (latest), macOS 27.0 (26A5421a), M5 Max,
`coreai-core==1.0.0b2`.
### Reproduction
```python
import torch, torch.nn as nn
from coreai_torch import TorchConverter, get_decomp_table
class M(nn.Module):
def __init__(self):
super().__init__()
self.p = nn.AvgPool2d(2, 2, 1, ceil_mode=True) # ceil_mode=True is the trigger
def forward(self, x):
return self.p(x)
ep = torch.export.export(M().eval(), args=(torch.randn(1, 3, 32, 32),))
ep = ep.run_decompositions(dict(get_decomp_table()))
TorchConverter().add_exported_program(ep, input_names=["x"], output_names=["out"]).to_coreai()
# IndexError: tuple index out of range
```
`torch.nn.functional.avg_pool2d(x, 2, 2, 1, True)` fails identically.
### Observed vs expected
| Construction | node `len(args)` | 0.4.1 | 0.4.2 |
|---|---|---|---|
| `nn.AvgPool2d(2, 2, 1, ceil_mode=True)` | 5 | `IndexError` | `IndexError` |
| `F.avg_pool2d(x, 2, 2, 1, True)` | 5 | `IndexError` | `IndexError` |
| `F.avg_pool2d(x, 2, 2, 1, True, False)` | 6 | converts | converts |
| `F.avg_pool2d(x, 2, 2, 0, False)` (defaults) | 3 | converts | converts |
Expected: all four convert. `count_include_pad` should default to `True` when absent.
### Root cause
The signature is
`aten.avg_pool2d(input, kernel_size, stride, padding, ceil_mode, count_include_pad, divisor_override)`,
so `ceil_mode` is index 4 and `count_include_pad` is index 5.
The two adjacent reads are:
```python
ceil_mode = (
node.args[4] if len(node.args) > 4 and node.args[4] is not None else False
)
count_include_pad = (
node.args[5] if len(node.args) > 4 and node.args[4] is not None else True
)
```
The `ceil_mode` line is correct. The `count_include_pad` line reads index **5** while still
guarding index **4** — the guard was not updated when the line was copied. With exactly five
args the guard is satisfied (`len > 4`, `args[4] is not None`) and the read runs off the end.
### Why it stays hidden
With a default `ceil_mode`, `torch.export` normalises the trailing arguments away and the node
carries only 3 args. The guard then fails cleanly and the `else True` branch happens to be
correct. The defect only surfaces when `ceil_mode` is explicitly non-default, which keeps element
4 in the graph while element 5 is still absent — so ordinary `AvgPool2d` usage converts fine and
`ceil_mode=True` does not.
### Suggested fix
```python
count_include_pad = (
node.args[5] if len(node.args) > 5 and node.args[5] is not None else True
)
```
### Workaround
Pass `count_include_pad` explicitly so the node carries six arguments:
```python
F.avg_pool2d(x, kernel_size, stride, padding, ceil_mode, count_include_pad)
```
### Note for downstream consumers
LibreYOLO carries a monkey-patch for this, scoped to
`_AFFECTED_COREAI_TORCH_VERSIONS = {"0.4.1"}`, which declines silently on any other version.
Since the defect is still present in 0.4.2, that shim now no-ops and the failure returns with no
diagnostic. Anyone pinning a version-scoped workaround for this will want to widen it until a fix
lands.
Contributor guide
Assessment
This issue has not been assessed yet.