Use lax.div in avg_pool so integer inputs stay integer

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),

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,21 +2258,22 @@ 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,
       kernel_size,
       strides,
       padding,
     )
+  y = jax.lax.div(y, jnp.broadcast_to(divisor, y.shape))
   return y.astype(inputs.dtype)