From 6a6e094d3b23d3afcb66c681ec00fa517be67268 Mon Sep 17 00:00:00 2001
From: Matthijs Douze <matthijs@fb.com>
Date: Thu, 28 Sep 2023 07:40:38 -0700
Subject: [PATCH] Relax IVFFlatDedup test

Summary:
This diff relaxes some IVFFlatDedup tests where distances are slighlty different over runs.
Should fix

https://app.circleci.com/pipelines/github/facebookresearch/faiss/4709/workflows/8c8213bf-8fe0-4c4e-9a7d-991f44bf1010/jobs/25551

https://app.circleci.com/pipelines/github/facebookresearch/faiss/4709/workflows/8c8213bf-8fe0-4c4e-9a7d-991f44bf1010/jobs/25547

Differential Revision: D49732349
---
 contrib/evaluation.py         |  3 ++-
 tests/test_index_composite.py | 23 ++++-------------------
 2 files changed, 6 insertions(+), 20 deletions(-)

diff --git a/contrib/evaluation.py b/contrib/evaluation.py
index 9c762b6081..9a19a34ad6 100644
--- a/contrib/evaluation.py
+++ b/contrib/evaluation.py
@@ -227,7 +227,8 @@ def compute_PR_for(q):
 # They are intended for use in tests
 
 def check_ref_knn_with_draws(Dref, Iref, Dnew, Inew):
-    """ test that knn search results are identical, raise if not """
+    """ test that knn search results are identical, with possible ties.
+    Raise if not. """
     np.testing.assert_array_almost_equal(Dref, Dnew, decimal=5)
     # here we have to be careful because of draws
     testcase = unittest.TestCase()   # because it makes nice error messages
diff --git a/tests/test_index_composite.py b/tests/test_index_composite.py
index d4f99b92d0..81a00cb938 100644
--- a/tests/test_index_composite.py
+++ b/tests/test_index_composite.py
@@ -17,7 +17,7 @@
 from common_faiss_tests import get_dataset_2
 from faiss.contrib.datasets import SyntheticDataset
 from faiss.contrib.inspect_tools import make_LinearTransform_matrix
-
+from faiss.contrib.evaluation import check_ref_knn_with_draws
 
 class TestRemoveFastScan(unittest.TestCase):
     def do_test(self, ntotal, removed):
@@ -430,12 +430,6 @@ def test_mmappedIO_pretrans(self):
 
 class TestIVFFlatDedup(unittest.TestCase):
 
-    def normalize_res(self, D, I):
-        dmax = D[-1]
-        res = [(d, i) for d, i in zip(D, I) if d < dmax]
-        res.sort()
-        return res
-
     def test_dedup(self):
         d = 10
         nb = 1000
@@ -471,10 +465,7 @@ def test_dedup(self):
         Dref, Iref = index_ref.search(xq, 20)
         Dnew, Inew = index_new.search(xq, 20)
 
-        for i in range(nq):
-            ref = self.normalize_res(Dref[i], Iref[i])
-            new = self.normalize_res(Dnew[i], Inew[i])
-            assert ref == new
+        check_ref_knn_with_draws(Dref, Iref, Dnew, Inew)
 
         # test I/O
         fd, tmpfile = tempfile.mkstemp()
@@ -487,10 +478,7 @@ def test_dedup(self):
                 os.unlink(tmpfile)
         Dst, Ist = index_st.search(xq, 20)
 
-        for i in range(nq):
-            new = self.normalize_res(Dnew[i], Inew[i])
-            st = self.normalize_res(Dst[i], Ist[i])
-            assert st == new
+        check_ref_knn_with_draws(Dnew, Inew, Dst, Ist)
 
         # test remove
         toremove = np.hstack((np.arange(3, 1000, 5), np.arange(850, 950)))
@@ -501,10 +489,7 @@ def test_dedup(self):
         Dref, Iref = index_ref.search(xq, 20)
         Dnew, Inew = index_new.search(xq, 20)
 
-        for i in range(nq):
-            ref = self.normalize_res(Dref[i], Iref[i])
-            new = self.normalize_res(Dnew[i], Inew[i])
-            assert ref == new
+        check_ref_knn_with_draws(Dref, Iref, Dnew, Inew)
 
 
 class TestSerialize(unittest.TestCase):