mirror of
https://github.com/facebookresearch/faiss.git
synced 2026-10-12 02:13:36 +00:00
Alibek Zhakubayev
55990eb91f
Summary: Pull Request resolved: https://github.com/facebookresearch/faiss/pull/5442 The RaBitQ utility tests landed as C++ (`test_rabitq_utils.cpp`) in D112153740, but faiss tests are primarily written in Python because they exercise the full SWIG binding stack. This diff ports them to Python with identical cases and no behavior change. The tested helpers (`round_nonnegative_to_uint8`, `round_nonnegative_to_uint16`, `round_clamped_to_uint8` in `faiss::rabitq_utils`, and `LinearTransform::apply_noalloc`) are all reachable from Python via the existing SWIG bindings, so no bindings changes are needed. The C++ test was registered only in `tests/CMakeLists.txt` and was never built by Buck; the Python port is added to `py_tests` in `tests/BUCK`, so it now runs under Buck as well. Changes: - Delete `tests/test_rabitq_utils.cpp` and remove it from `tests/CMakeLists.txt`. - Add `tests/test_rabitq_utils.py` with the same cases: the fixed-value assertions for all three round helpers, the 10000-sample +/-1 property tests, and the `LinearTransform` single-vector vs batch check for both bias settings. - Register `test_rabitq_utils` in `py_tests` in `tests/BUCK`. The random property tests use numpy RNG rather than std::mt19937 (the two RNG engines cannot be made bit-identical across languages), but the ranges, sample counts, and +/-1 tolerance are unchanged. Reviewed By: mnorris11 Differential Revision: D112559549 fbshipit-source-id: 9be93dc33e94477c6c5919582d3a8b4774e476c1
115 lines
4.0 KiB
Python
115 lines
4.0 KiB
Python
# Copyright (c) Meta Platforms, Inc. and affiliates.
|
|
#
|
|
# This source code is licensed under the MIT license found in the
|
|
# LICENSE file in the root directory of this source tree.
|
|
|
|
import math
|
|
import unittest
|
|
|
|
import faiss
|
|
|
|
import numpy as np
|
|
|
|
|
|
def round_half_away(x):
|
|
# Matches C++ std::round (round half away from zero), unlike Python round()
|
|
# which rounds half to even.
|
|
return math.floor(x + 0.5) if x >= 0 else math.ceil(x - 0.5)
|
|
|
|
|
|
def round_clamped_ref(x, max_code):
|
|
# Mirrors std::round (half away from zero) then clamp to [0, max_code].
|
|
y = round_half_away(x)
|
|
y = max(0.0, min(y, float(max_code)))
|
|
return int(y)
|
|
|
|
|
|
def check_linear_transform_batch_match(test, have_bias):
|
|
d_in = 7
|
|
d_out = 5
|
|
transform = faiss.LinearTransform(d_in, d_out, have_bias)
|
|
transform.is_trained = True
|
|
|
|
a = np.sin(np.arange(d_in * d_out, dtype=np.float32) * np.float32(0.37))
|
|
faiss.copy_array_to_vector(a, transform.A)
|
|
|
|
if have_bias:
|
|
b = np.arange(d_out, dtype=np.float32)
|
|
b = np.float32(0.25) * b - np.float32(0.5)
|
|
faiss.copy_array_to_vector(b, transform.b)
|
|
|
|
x = np.cos(np.arange(d_in, dtype=np.float32) * np.float32(0.19))
|
|
|
|
one_out = np.empty(d_out, dtype="float32")
|
|
transform.apply_noalloc(1, faiss.swig_ptr(x), faiss.swig_ptr(one_out))
|
|
|
|
batch_x = np.tile(x, 4)
|
|
batch_out = np.empty(4 * d_out, dtype="float32")
|
|
transform.apply_noalloc(
|
|
4, faiss.swig_ptr(batch_x), faiss.swig_ptr(batch_out)
|
|
)
|
|
|
|
for i in range(d_out):
|
|
test.assertAlmostEqual(one_out[i], batch_out[i], delta=1e-5)
|
|
|
|
|
|
class RaBitQUtilsTest(unittest.TestCase):
|
|
def test_round_nonnegative_to_uint8(self):
|
|
self.assertEqual(faiss.round_nonnegative_to_uint8(0.0), 0)
|
|
self.assertEqual(faiss.round_nonnegative_to_uint8(0.49), 0)
|
|
self.assertEqual(faiss.round_nonnegative_to_uint8(0.5), 1)
|
|
self.assertEqual(faiss.round_nonnegative_to_uint8(1.5), 2)
|
|
self.assertEqual(faiss.round_nonnegative_to_uint8(254.5), 255)
|
|
self.assertEqual(faiss.round_nonnegative_to_uint8(255.0), 255)
|
|
|
|
below_half = float(np.nextafter(np.float32(0.5), np.float32(0.0)))
|
|
self.assertLessEqual(
|
|
abs(
|
|
int(faiss.round_nonnegative_to_uint8(below_half))
|
|
- round_half_away(below_half)
|
|
),
|
|
1,
|
|
)
|
|
|
|
rng = np.random.RandomState(123)
|
|
for x in rng.uniform(0.0, 255.49, size=10000).astype(np.float32):
|
|
x = float(x)
|
|
self.assertLessEqual(
|
|
abs(
|
|
int(faiss.round_nonnegative_to_uint8(x))
|
|
- round_half_away(x)
|
|
),
|
|
1,
|
|
)
|
|
|
|
def test_round_nonnegative_to_uint16(self):
|
|
self.assertEqual(faiss.round_nonnegative_to_uint16(0.0), 0)
|
|
self.assertEqual(faiss.round_nonnegative_to_uint16(0.5), 1)
|
|
self.assertEqual(faiss.round_nonnegative_to_uint16(65534.5), 65535)
|
|
self.assertEqual(faiss.round_nonnegative_to_uint16(65535.0), 65535)
|
|
|
|
def test_round_clamped_to_uint8(self):
|
|
self.assertEqual(faiss.round_clamped_to_uint8(-1.0, 15), 0)
|
|
self.assertEqual(faiss.round_clamped_to_uint8(0.49, 15), 0)
|
|
self.assertEqual(faiss.round_clamped_to_uint8(0.5, 15), 1)
|
|
self.assertEqual(faiss.round_clamped_to_uint8(14.5, 15), 15)
|
|
self.assertEqual(faiss.round_clamped_to_uint8(15.4, 15), 15)
|
|
self.assertEqual(faiss.round_clamped_to_uint8(16.0, 15), 15)
|
|
|
|
rng = np.random.RandomState(456)
|
|
for x in rng.uniform(-4.0, 260.0, size=10000).astype(np.float32):
|
|
x = float(x)
|
|
self.assertLessEqual(
|
|
abs(
|
|
int(faiss.round_clamped_to_uint8(x, 255))
|
|
- round_clamped_ref(x, 255)
|
|
),
|
|
1,
|
|
)
|
|
|
|
|
|
class LinearTransformTest(unittest.TestCase):
|
|
def test_single_vector_matches_batch_path(self):
|
|
check_linear_transform_batch_match(self, False)
|
|
check_linear_transform_batch_match(self, True)
|