Fix BruteForce serialize test (#1568)

`test_serialization.py::test_save_load_brute_force` can sometimes fail due to the ordering of the neighbors. To fix this we can sort the neighbors before comparison, and if there's a small discrepancy due to two neighbors being at the same distance on the `k` cutoff we check that the recall is over 99.8%.

I ran this test on a loop for 60k iterations without it failing.

Authors:
  - Micka (https://github.com/lowener)

Approvers:
  - Gil Forsyth (https://github.com/gforsyth)
  - Divye Gala (https://github.com/divyegala)
  - Corey J. Nolet (https://github.com/cjnolet)

URL: #1568

Author: lowener

python/cuvs/cuvs/tests/test_serialization.py

@@ -1,4 +1,4 @@
-# SPDX-FileCopyrightText: Copyright (c) 2024, NVIDIA CORPORATION.
+# SPDX-FileCopyrightText: Copyright (c) 2024-2025, NVIDIA CORPORATION.
 # SPDX-License-Identifier: Apache-2.0
 #
 
@@ -7,7 +7,7 @@
 from pylibraft.common import device_ndarray
 
 from cuvs.neighbors import brute_force, cagra, ivf_flat, ivf_pq
-from cuvs.tests.ann_utils import generate_data
+from cuvs.tests.ann_utils import calc_recall, generate_data
 
 
 @pytest.mark.parametrize("dtype", [np.float32, np.int8, np.ubyte])
@@ -77,5 +77,17 @@ def run_save_load(ann_module, dtype):
     neighbors2 = neighbors_dev.copy_to_host()
     dist2 = distance_dev.copy_to_host()
 
-    assert np.all(neighbors == neighbors2)
     assert np.allclose(dist, dist2, rtol=1e-6)
+
+    # Sort the neighbors to avoid ordering issues
+    sorted_neighbors = np.argsort(neighbors, axis=-1)
+    sorted_neighbors2 = np.argsort(neighbors2, axis=-1)
+    neighbors = np.take_along_axis(neighbors, sorted_neighbors, axis=-1)
+    neighbors2 = np.take_along_axis(neighbors2, sorted_neighbors2, axis=-1)
+    all_match = np.all(neighbors == neighbors2)
+    # If the neighbors are not the same, there might be a cutoff between the k
+    # and k+1 neighbors at the same distance.
+    # Calculate that the recall is at least 99.8%
+    if not all_match:
+        recall = calc_recall(neighbors, neighbors2)
+        assert recall >= 0.998