apple / apple/coreai-torch

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

Open
#73 0 comments 0 reactions 1 assignee Claimed by @jakesabathia2 View on GitHub
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

Open the contributing guide

Assessment

This issue has not been assessed yet.

Get new issues in your inbox

A short digest of beginner-friendly GitHub issues.