[john-rocky]

Summary

PyTorch accepts avg_pool2d on an input without a batch dim, and nn.AvgPool1d on [C, L] reaches avg_pool2d as [C, 1, L]. replace_avg_pool2d assumed a rank-4 input and raised IndexError at x.type.shape[3].

For a rank-3 input, the lowering now adds a batch dim of 1 before the composite and removes it after. Rank-4 inputs lower as before.

git merge-tree against the heads of #80, #81 and #123, which also change this function, reports no conflict. Unbatched AvgPool3d, MaxPool1d / MaxPool2d and AdaptiveAvgPool1d / AdaptiveAvgPool2d still fail to convert, on main and with this change; this PR covers only #144.

Testing

test_avg_pool2d_unbatched runs nn.AvgPool1d(2) and nn.AvgPool1d(3, stride=2, padding=1) on [4, 8] and nn.AvgPool2d(2) on [3, 8, 8], static and with a dynamic channel dim, against PyTorch. All six raise IndexError before the fix and pass after it.

uv run pytest tests/ops/test_ops.py -k avg_pool             # 102 passed
uv run pytest tests/ops/test_ops_ir.py -k TestAvgPool2dIR   # 2 passed
uv run pytest tests/ -n 4                                   # 6560 passed; 3 failed in tests/debugging/, as on main
uv run --extra dev ruff check .                             # All checks passed
uv run --extra dev ruff format --check .                    # all files formatted

Fixes #144