Skip to content
Merged
Show file tree
Hide file tree
Changes from all commits
Commits
File filter

Filter by extension

Filter by extension

Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
2 changes: 2 additions & 0 deletions docs/semantics.md
Original file line number Diff line number Diff line change
Expand Up @@ -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
Expand Down
15 changes: 13 additions & 2 deletions src/array_vectorize/runtime/arith.py
Original file line number Diff line number Diff line change
@@ -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
Expand Down Expand Up @@ -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)
Expand Down
25 changes: 25 additions & 0 deletions tests/backends/test_jax_torch_compilation.py
Original file line number Diff line number Diff line change
@@ -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``
Expand Down Expand Up @@ -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))
Loading