From 7623b13b5effee010c32fd6640be9501f9a6f05c Mon Sep 17 00:00:00 2001 From: Benjamin Sepanski Date: Thu, 3 Sep 2026 11:01:04 -0500 Subject: [PATCH] Use lax.div in avg_pool so integer inputs stay integer `_aten_avg_pool` divided the window sums with `/`, which jnp promotes to floating point for integer inputs, so an integer avg_pool lowered to convert -> divide(f32) -> convert. Use `jax.lax.div` instead: for integer dtypes it is truncating division, matching torch's semantics for int64 avg_pool, and the lowering no longer introduces intermediate float values. Add int64 tests across the three divisor branches (count_include_pad with ceil_mode, count_include_pad=False, divisor_override) plus a check that the lowered StableHLO contains no f32. --- test/test_core_aten_ops.py | 43 ++++++++++++++++++++++++++++++++++++++ torchax/ops/jaten.py | 9 +++++--- 2 files changed, 49 insertions(+), 3 deletions(-) diff --git a/test/test_core_aten_ops.py b/test/test_core_aten_ops.py index 60f5ac0..7056a6d 100644 --- a/test/test_core_aten_ops.py +++ b/test/test_core_aten_ops.py @@ -12,13 +12,17 @@ # See the License for the specific language governing permissions and # limitations under the License. +import functools import math import unittest +import jax +import jax.numpy as jnp import torch from torch.utils import _pytree as pytree from torchax import tensor +from torchax.ops import jaten from . import base_test_util @@ -780,6 +784,45 @@ def test_aten_avg_pool3d_0(self): kwargs = {} run_export_and_compare(self, torch.ops.aten.avg_pool3d, args, kwargs) + def test_aten_avg_pool2d_int(self): + # Negative values exercise truncating (toward-zero) integer division. + args = ( + torch.randint(-20, 20, (1, 3, 7, 9)).to(torch.int64), + [3, 3], + [2, 2], + [1, 1], + ) + for kwargs in ( + {"ceil_mode": True, "count_include_pad": True}, + {"count_include_pad": False}, + {"divisor_override": 4}, + ): + with self.subTest(kwargs=kwargs): + run_export_and_compare(self, torch.ops.aten.avg_pool2d, args, kwargs) + + def test_aten_avg_pool3d_int(self): + args = ( + torch.randint(-20, 20, (1, 2, 6, 7, 8)).to(torch.int64), + [2, 3, 2], + [2, 2, 2], + [1, 1, 0], + ) + kwargs = {"count_include_pad": False} + run_export_and_compare(self, torch.ops.aten.avg_pool3d, args, kwargs) + + def test_aten_avg_pool2d_int_lowers_without_float(self): + x = jnp.zeros((1, 3, 8, 8), dtype=jnp.int32) + for kwargs in ( + {}, + {"count_include_pad": False}, + {"divisor_override": 3}, + ): + f = functools.partial( + jaten._aten_avg_pool, kernel_size=(3, 3), padding=1, **kwargs + ) + hlo = jax.jit(f).lower(x).as_text() + self.assertNotIn("f32", hlo, msg=f"kwargs={kwargs}:\n{hlo}") + def test_aten_bitwise_and_Scalar_0(self): args = ( torch.randint(0, 10, (10, 10)).to(torch.int32), diff --git a/torchax/ops/jaten.py b/torchax/ops/jaten.py index cae98e3..cbc915b 100644 --- a/torchax/ops/jaten.py +++ b/torchax/ops/jaten.py @@ -2234,8 +2234,10 @@ def _aten_avg_pool( ) y = pool(inputs, 0.0, jax.lax.add, kernel_size, strides, padding) + # lax.div (not `/`) keeps integer inputs integer: truncating division like + # torch, with no round trip through floating point. if divisor_override is not None: - y = y / jnp.array(divisor_override, y.dtype) + divisor = jnp.array(divisor_override, y.dtype) elif count_include_pad: div_shape = list(y.shape) div_by = jnp.ones(div_shape, y.dtype) * np.prod(kernel_size) @@ -2256,14 +2258,14 @@ def _aten_avg_pool( idx[j + offset] = -1 div_by = div_by.at[tuple(idx)].set(np.prod(new_kernel_size)) - y = y / div_by + divisor = div_by else: div_shape = list(inputs.shape) div_shape[num_batch_dims] = 1 div_shape = tuple(div_shape) if len(div_shape) - 2 == len(kernel_size): div_shape = (1,) + div_shape[1:] - y = y / pool( + divisor = pool( jnp.ones(div_shape, y.dtype), jnp.array(0.0, y.dtype), jax.lax.add, @@ -2271,6 +2273,7 @@ def _aten_avg_pool( strides, padding, ) + y = jax.lax.div(y, jnp.broadcast_to(divisor, y.shape)) return y.astype(inputs.dtype)