Use lax.div in avg_pool so integer inputs stay integer

[benSepanski]

Fixes #105.

Summary

_aten_avg_pool in torchax/ops/jaten.py divided the window sums with /. For integer inputs, jnp promotes that to floating point, so an integer avg_pool lowered to convert(i32->f32) → divide → convert(f32->i32).

This PR switches the division to jax.lax.div, which for integer dtypes is truncating (toward-zero) division. That matches torch's behavior for int64 avg_pool{1,2,3}d (e.g. -7 / 4 == -1), and the lowering no longer introduces intermediate floating-point values. Floating-point inputs are unaffected: lax.div is a plain elementwise divide there, same as /.

Before (int32 input, 3x3 kernel):

%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>

After:

%4 = stablehlo.divide %1, %3 : tensor<1x1x2x2xi32>

Tests

• test_aten_avg_pool2d_int: int64 inputs with negative values (to exercise truncation) across all three divisor branches: count_include_pad=True with ceil_mode, count_include_pad=False, and divisor_override.
• test_aten_avg_pool3d_int: int64 3d case with count_include_pad=False.
• test_aten_avg_pool2d_int_lowers_without_float: asserts the lowered StableHLO for an int32 input contains no f32 in each branch. This test fails on main.

Ran pytest test/test_core_aten_ops.py (all pass) and the avg_pool OpInfo tests in test/test_ops.py (all pass). The max_pool{2,3}d int64 OpInfo tests fail on main as well and are unrelated.

[google-cla]

Thanks for your pull request! It looks like this may be your first contribution to a Google open source project. Before we can look at your pull request, you'll need to sign a Contributor License Agreement (CLA).

View this failed invocation of the CLA check for more information.

For the most up to date status, view the checks section at the bottom of the pull request.