avg_pool on integer inputs lowers through float32
- Dominant language
- Python
- Stars
- 243
- Forks
- 38
- Avg merge
- 15h 32m
- Merged PRs (30d)
- 2
Description
## Problem
`torch.ops.aten.avg_pool{1,2,3}d` on an integer tensor currently lowers to convert-to-float, divide, convert-back. `_aten_avg_pool` in `torchax/ops/jaten.py` divides the window sums with `/`, and `jnp` promotes integer `/` to floating point.
For an int32 input with a 3x3 kernel the StableHLO looks like:
```mlir
%1 = "stablehlo.reduce_window"(...) : (tensor<1x1x4x4xi32>, tensor) -> tensor<1x1x2x2xi32>
...
%5 = stablehlo.convert %1 : (tensor<1x1x2x2xi32>) -> tensor<1x1x2x2xf32>
%6 = stablehlo.convert %4 : (tensor<1x1x2x2xi32>) -> tensor<1x1x2x2xf32>
%7 = stablehlo.divide %5, %6 : tensor<1x1x2x2xf32>
%8 = stablehlo.convert %7 : (tensor<1x1x2x2xf32>) -> tensor<1x1x2x2xi32>
```
Two consequences:
- The lowered program contains intermediate float values for an all-integer computation. That matters for targets that want integer-only graphs.
- Results can differ from torch once values exceed f32's 24-bit mantissa:
```python
x = torch.full((1, 1, 2, 2), 2**24 + 1, dtype=torch.int64)
torch.nn.functional.avg_pool2d(x, 1) # 16777217 (exact)
# same call under torchax on main # 16777216
```
## Torch semantics
Torch's CPU kernel supports `avg_pool{2,3}d` on `int64` and computes `sum / divisor` with C++ integer division, which truncates toward zero. Other integer dtypes raise "not implemented", so int64 is the only integer case with defined behavior.
## Proposal
Replace the `/` in `_aten_avg_pool` with `jax.lax.div`. For integer dtypes `lax.div` is truncating division, matching torch, and the lowering becomes a single integer `stablehlo.divide`. For float dtypes `lax.div` is the same elementwise divide as `/`, so float behavior is unchanged.
I have a draft PR implementing this with tests: #104
Contributor guide
Research direction
Start in torchax/ops/jaten.py at _aten_avg_pool, then review draft PR #104 and its tests. Verify integer pooling preserves torch's int64 truncation behavior without intermediate float conversion, while float pooling remains unchanged; the relevant tests are included in the draft PR.
Written by the indexing model from the issue text.
Assessment
- Tech stack
- python, pytorch
- Domain
- compilers, machine-learning
- Issue type
- Bug
- Difficulty
- 2/5
- Estimated time
- 1-3 hours
- Activity status
- Stale
- Clarity
- Clearly specified
- Newbie friendliness
- 35/100