diff --git a/docs/semantics.md b/docs/semantics.md index 96c793a..371e542 100644 --- a/docs/semantics.md +++ b/docs/semantics.md @@ -48,6 +48,8 @@ call and checked at runtime on the actual values: True division converts integer and boolean operands to floating arrays, including inputs whose dtype is known only at call time. Integer ratios round once to float64, without first rounding large integer operands. +This includes the signed 64-bit minimum in either operand under JIT +compilation, with the same rounding and sign of zero as scalar Python. This uses a fixed bit loop, which costs more at call time and first JIT compilation than floating-point division. Negative integer literals (also when assigned to local names) use the exact unsigned diff --git a/src/array_vectorize/runtime/arith.py b/src/array_vectorize/runtime/arith.py index ccd2e94..9e52f4a 100644 --- a/src/array_vectorize/runtime/arith.py +++ b/src/array_vectorize/runtime/arith.py @@ -1,3 +1,7 @@ +# Copyright (c) 2026 grAItools +# SPDX-License-Identifier: BSD-3-Clause +# See LICENSE for the full license text. + """Exact-semantics arithmetic helper (from _runtime). ``_vec_arith`` and its uint64-exactness primitives are injected into @@ -133,8 +137,15 @@ def _integer_true_divide(xp: Any, left: Any, right: Any) -> Any: b = xp.astype(right_array, xp.int64) l_neg = a < 0 if "uint" not in str(left_array.dtype) else xp.asarray(False) r_neg = b < 0 if "uint" not in str(right_array.dtype) else xp.asarray(False) - a = xp.where(l_neg, -a, a) - b = xp.where(r_neg, -b, b) + # Preserve the minimum's unsigned magnitude bit pattern without negating + # it: scalar compiled kernels may treat signed overflow as undefined. + # Both where branches evaluate (design D1), so mask before negation. + a_min = a == -(2**63) + b_min = b == -(2**63) + safe_a = xp.where(a_min, xp.asarray(0, dtype=xp.int64), a) + safe_b = xp.where(b_min, xp.asarray(0, dtype=xp.int64), b) + a = xp.where(a_min, a, xp.where(l_neg, -safe_a, a)) + b = xp.where(b_min, b, xp.where(r_neg, -safe_b, b)) zero_divisor = b == 0 b = xp.where(zero_divisor, xp.asarray(1, dtype=xp.int64), b) leading_position = _integer_exponent(xp, a) - _integer_exponent(xp, b) diff --git a/tests/backends/test_jax_torch_compilation.py b/tests/backends/test_jax_torch_compilation.py index b8cab11..84fba93 100644 --- a/tests/backends/test_jax_torch_compilation.py +++ b/tests/backends/test_jax_torch_compilation.py @@ -1,3 +1,7 @@ +# Copyright (c) 2026 grAItools +# SPDX-License-Identifier: BSD-3-Clause +# See LICENSE for the full license text. + """jax.jit / torch.compile compatibility + backend-detection regression tests. JAX and PyTorch trace and introspect the Python they compile — ``jax.jit`` @@ -282,3 +286,24 @@ def test_compiled_integer_division_preserves_python_rounding(backend: str) -> No expected = [a / b for a, b in zip(left, right, strict=True)] np.testing.assert_array_equal(np.asarray(got), expected) np.testing.assert_array_equal(np.signbit(np.asarray(got)), np.signbit(expected)) + + +@pytest.mark.timeout(600) +def test_scalar_compiled_integer_division_preserves_minimum_magnitudes() -> None: + torch = pytest.importorskip("torch") + import array_api_compat.torch + import torch._inductor.config as inductor_config + + scalar = support.make_fn(" return x / y", defaults="x, y") + minimum = -(2**63) + left = [minimum, minimum, minimum, 1, -1, 0, 2**63 - 1] + right = [-1, 1, minimum, minimum, minimum, minimum, 3] + vec = array_vectorize.vectorize(scalar, namespace=array_api_compat.torch) + # Force the scalar C++ kernel even on hosts with vector instructions. + with inductor_config.patch({"cpp.simdlen": 1}): + got = torch.compile(vec)( + torch.tensor(left, dtype=torch.int64), torch.tensor(right, dtype=torch.int64) + ) + expected = [a / b for a, b in zip(left, right, strict=True)] + np.testing.assert_array_equal(np.asarray(got), expected) + np.testing.assert_array_equal(np.signbit(np.asarray(got)), np.signbit(expected))