From 710d5dbf32ee5044bfe26e48160ab1bf2541bc04 Mon Sep 17 00:00:00 2001 From: stephantul Date: Fri, 2 Oct 2026 08:49:25 +0200 Subject: [PATCH] fix: adaptive epsilon for dtypes --- tests/test_utils.py | 11 +++++++++++ vicinity/utils.py | 5 ++++- 2 files changed, 15 insertions(+), 1 deletion(-) diff --git a/tests/test_utils.py b/tests/test_utils.py index 6abd937..3f5f083 100644 --- a/tests/test_utils.py +++ b/tests/test_utils.py @@ -43,3 +43,14 @@ def test_normalize_or_copy() -> None: result_zero = normalize_or_copy(zero_vectors) assert_array_equal(result_zero, zero_vectors) assert result_zero is zero_vectors, "Should return the original array" + + +def test_normalize_or_copy_dtypes() -> None: + """Test that normalized vectors are recognized for every float dtype.""" + rng = np.random.default_rng(42) + for dtype in (np.float16, np.float32, np.float64): + vectors = normalize(rng.standard_normal((1000, 384)).astype(dtype)) + assert normalize_or_copy(vectors) is vectors, f"Should return the original array for {dtype}" + + scaled = vectors * dtype(1.01) + assert normalize_or_copy(scaled) is not scaled, f"Should return a new array for {dtype}" diff --git a/vicinity/utils.py b/vicinity/utils.py index b46de26..2513b73 100644 --- a/vicinity/utils.py +++ b/vicinity/utils.py @@ -49,7 +49,10 @@ def normalize_or_copy(vectors: npt.NDArray) -> npt.NDArray: Otherwise, the vectors are normalized, and a new array is returned. """ norms = np.linalg.norm(vectors, axis=-1) - all_unit_length = np.allclose(norms[norms != 0], 1) + rtol = 1e-5 + if np.issubdtype(vectors.dtype, np.floating): + rtol = max(rtol, 4 * float(np.finfo(vectors.dtype).eps)) + all_unit_length = np.allclose(norms[norms != 0], 1, rtol=rtol) if all_unit_length: return vectors return normalize(vectors, norms)