diff --git a/contrib/evaluation.py b/contrib/evaluation.py index 927c88e66..6ed969b81 100644 --- a/contrib/evaluation.py +++ b/contrib/evaluation.py @@ -244,10 +244,10 @@ def _cluster_tables_with_tolerance(tab1, tab2, thr): return idx1, idx2 -def check_ref_knn_with_draws(Dref, Iref, Dnew, Inew, rtol=1e-5): +def check_ref_knn_with_draws(Dref, Iref, Dnew, Inew, rtol=1e-5, atol=0): """test that knn search results are identical, with possible ties. Raise if not.""" - np.testing.assert_allclose(Dref, Dnew, rtol=rtol) + np.testing.assert_allclose(Dref, Dnew, rtol=rtol, atol=atol) # here we have to be careful because of draws testcase = unittest.TestCase() # because it makes nice error messages for i in range(len(Iref)): @@ -255,7 +255,7 @@ def check_ref_knn_with_draws(Dref, Iref, Dnew, Inew, rtol=1e-5): continue # otherwise collect elements per distance - r = rtol * Dref[i].max() + r = rtol * Dref[i].max() + atol DrefC, DnewC = _cluster_tables_with_tolerance(Dref[i], Dnew[i], r) diff --git a/tests/test_ivf_flat_panorama.py b/tests/test_ivf_flat_panorama.py index 638f69f82..4c250b04d 100644 --- a/tests/test_ivf_flat_panorama.py +++ b/tests/test_ivf_flat_panorama.py @@ -105,7 +105,7 @@ class TestIndexIVFFlatPanorama(unittest.TestCase): err_msg="Distances mismatch", ) check_ref_knn_with_draws( - D_regular, I_regular, D_regular, I_panorama, rtol=rtol + D_regular, I_regular, D_panorama, I_panorama, rtol=rtol, atol=atol ) def assert_range_results_equal( diff --git a/tests/test_refine_panorama.py b/tests/test_refine_panorama.py index a8555a7fa..889b511ac 100644 --- a/tests/test_refine_panorama.py +++ b/tests/test_refine_panorama.py @@ -91,7 +91,7 @@ class TestIndexRefinePanorama(unittest.TestCase): err_msg="Distances mismatch", ) check_ref_knn_with_draws( - D_regular, I_regular, D_regular, I_panorama, rtol=rtol + D_regular, I_regular, D_panorama, I_panorama, rtol=rtol, atol=atol ) def test_refine_panorama_correctness(self):