google / google/torchax

avg_pool on integer inputs lowers through float32

Open
#105 0 comments 0 reactions 0 assignees View on GitHub
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

Open the contributing 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

Get new issues in your inbox

A short digest of beginner-friendly GitHub issues.