# 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. from __future__ import absolute_import, division, print_function import os import tempfile import numpy as np import faiss import unittest from faiss.contrib.datasets import SyntheticDataset from common_faiss_tests import for_all_simd_levels, NoneSIMDLevel @for_all_simd_levels class TestScalarQuantizerEncodeDecode(unittest.TestCase): def setUp(self): self.d = 32 self.ds = SyntheticDataset(d=self.d, nt=0, nb=1000, nq=0, seed=42) self.xb = self.ds.get_database() def do_encode_decode(self, qtype, max_err): sq = faiss.ScalarQuantizer(self.d, qtype) sq.train(self.xb) codes = sq.compute_codes(self.xb) xb_decoded = sq.decode(codes) self.assertEqual(xb_decoded.shape, self.xb.shape) err = np.abs(self.xb - xb_decoded).max() self.assertLess(err, max_err) def test_4bit(self): self.do_encode_decode(faiss.ScalarQuantizer.QT_4bit, 0.1) def test_fp16(self): self.do_encode_decode(faiss.ScalarQuantizer.QT_fp16, 0.001) def test_8bit_uniform(self): self.do_encode_decode(faiss.ScalarQuantizer.QT_8bit_uniform, 0.01) def test_4bit_uniform(self): self.do_encode_decode(faiss.ScalarQuantizer.QT_4bit_uniform, 0.1) def test_tqmse_1bit(self): faiss.normalize_L2(self.xb) self.do_encode_decode(faiss.ScalarQuantizer.QT_1bit_tqmse, 1.0) def test_tqmse_2bit(self): faiss.normalize_L2(self.xb) self.do_encode_decode(faiss.ScalarQuantizer.QT_2bit_tqmse, 0.8) def test_tqmse_3bit(self): faiss.normalize_L2(self.xb) self.do_encode_decode(faiss.ScalarQuantizer.QT_3bit_tqmse, 0.6) def test_tqmse_4bit(self): faiss.normalize_L2(self.xb) self.do_encode_decode(faiss.ScalarQuantizer.QT_4bit_tqmse, 0.2) def test_tqmse_8bit(self): faiss.normalize_L2(self.xb) self.do_encode_decode(faiss.ScalarQuantizer.QT_8bit_tqmse, 0.1) def test_tqmse_preserves_norm(self): """tqmse encodes unit-norm vectors with no per-vector scale factor in the code, so reconstructions must keep unit norm. The Lloyd-Max table is optimal for unit *variance*, so it only lines up with unit-norm input once scaled by 1/sqrt(d) -- without that, every component lands in the innermost cells and the codebook collapses toward a sign quantizer, inflating norms by 15% at d=64 and 258% at d=768.""" for qtype, nbits in ( (faiss.ScalarQuantizer.QT_1bit_tqmse, 1), (faiss.ScalarQuantizer.QT_2bit_tqmse, 2), (faiss.ScalarQuantizer.QT_3bit_tqmse, 3), (faiss.ScalarQuantizer.QT_4bit_tqmse, 4), (faiss.ScalarQuantizer.QT_8bit_tqmse, 8), ): for d in (32, 64, 768): with self.subTest(nbits=nbits, d=d): rs = np.random.RandomState(123) x = rs.randn(200, d).astype("float32") faiss.normalize_L2(x) sq = faiss.ScalarQuantizer(d, qtype) sq.train(x) decoded = sq.decode(sq.compute_codes(x)) norms = np.linalg.norm(decoded, axis=1) # 1-bit cannot resolve magnitude at all (every component # reconstructs to +/-c), so it gets a looser band. tol = 0.25 if nbits == 1 else 0.1 self.assertLess(abs(np.median(norms) - 1.0), tol) # Guard the collapse directly: a codebook matched to the # data reconstructs most of its 2^nbits distinct levels. # This is what catches 8-bit, whose norm stays near 1 even # when collapsed (31 of 256 levels used at d=768). used = len(np.unique(decoded)) self.assertGreaterEqual(used, 2**nbits // 2) def test_codes_match_none(self): """SQ codes are integer; encode dispatch (sq-dispatch.h) must produce bit-identical output at every SIMD level. Catches drift in the per-ISA scalar-quantizer encode kernels.""" if not faiss.SIMDConfig.is_simd_level_available(faiss.SIMDLevel_NONE): self.skipTest("SIMDLevel.NONE not available") for qtype in ( faiss.ScalarQuantizer.QT_8bit, faiss.ScalarQuantizer.QT_4bit, faiss.ScalarQuantizer.QT_8bit_uniform, faiss.ScalarQuantizer.QT_4bit_uniform, faiss.ScalarQuantizer.QT_fp16, ): with self.subTest(qtype=qtype): sq = faiss.ScalarQuantizer(self.d, qtype) sq.train(self.xb) codes = sq.compute_codes(self.xb) with NoneSIMDLevel(): codes_none = sq.compute_codes(self.xb) np.testing.assert_array_equal(codes, codes_none) @for_all_simd_levels class TestScalarQuantizerSearch(unittest.TestCase): def setUp(self): self.d = 32 self.ds = SyntheticDataset(d=self.d, nt=0, nb=10000, nq=100, seed=42) self.xb = self.ds.get_database() self.xq = self.ds.get_queries() self.xb_unit = self.xb.copy() self.xq_unit = self.xq.copy() faiss.normalize_L2(self.xb_unit) faiss.normalize_L2(self.xq_unit) def do_search(self, factory_str, min_recall): index_gt = faiss.IndexFlatL2(self.d) index_gt.add(self.xb) _, I_gt = index_gt.search(self.xq, 10) index = faiss.index_factory(self.d, factory_str) index.train(self.xb) index.add(self.xb) _, I = index.search(self.xq, 10) recall = (I_gt[:, 0] == I[:, 0]).mean() self.assertGreater(recall, min_recall) def test_SQ4(self): self.do_search("SQ4", 0.5) def test_SQfp16(self): self.do_search("SQfp16", 0.99) def test_tqmse_search(self): index_gt = faiss.IndexFlatL2(self.d) index_gt.add(self.xb_unit) _, I_gt = index_gt.search(self.xq_unit, 10) recalls = {} for factory_str in ["L2norm,RR,SQtqmse4", "L2norm,RR,SQtqmse8"]: index = faiss.index_factory(self.d, factory_str) index.train(self.xb) index.add(self.xb) _, I = index.search(self.xq, 10) recalls[factory_str] = (I_gt[:, 0] == I[:, 0]).mean() self.assertGreater(recalls["L2norm,RR,SQtqmse4"], 0.2) self.assertGreaterEqual( recalls["L2norm,RR,SQtqmse8"], recalls["L2norm,RR,SQtqmse4"] ) @for_all_simd_levels class TestScalarQuantizerDistances(unittest.TestCase): def test_distance_matches_reconstruct(self): d = 32 ds = SyntheticDataset(d=d, nt=0, nb=100, nq=1, seed=42) x = ds.get_queries() xb = ds.get_database() index = faiss.index_factory(d, "SQ8") index.train(xb) index.add(xb) D, I = index.search(x, 10) for i in range(10): xr = index.reconstruct(int(I[0, i])) dist = ((x[0] - xr) ** 2).sum() self.assertAlmostEqual(D[0, i], dist, places=4) def test_tqmse_distance_matches_reconstruct(self): d = 32 ds = SyntheticDataset(d=d, nt=0, nb=100, nq=1, seed=42) x = ds.get_queries() xb = ds.get_database() faiss.normalize_L2(x) faiss.normalize_L2(xb) index = faiss.index_factory(d, "SQtqmse8") index.train(xb) index.add(xb) D, I = index.search(x, 10) for i in range(10): xr = index.reconstruct(int(I[0, i])) dist = ((x[0] - xr) ** 2).sum() self.assertAlmostEqual(D[0, i], dist, places=4) @for_all_simd_levels class TestScalarQuantizerEdgeCases(unittest.TestCase): def test_zero_vectors(self): d = 32 xb = np.zeros((100, d), dtype="float32") index = faiss.index_factory(d, "SQ8") index.train(xb) index.add(xb) D, _ = index.search(np.zeros((1, d), dtype="float32"), 10) self.assertTrue(np.allclose(D, 0, atol=1e-5)) def test_constant_vectors(self): d = 32 xb = np.ones((100, d), dtype="float32") * 0.5 index = faiss.index_factory(d, "SQ8") index.train(xb) index.add(xb) D, _ = index.search(np.ones((1, d), dtype="float32") * 0.5, 10) self.assertTrue(np.allclose(D, 0, atol=1e-3)) def test_extreme_dims(self): """Test extreme dimension values (very small and very large).""" for d in [1, 1024]: with self.subTest(d=d): ds = SyntheticDataset(d=d, nt=0, nb=100, nq=10, seed=42) xb = ds.get_database() xq = ds.get_queries() index = faiss.index_factory(d, "SQ8") index.train(xb) index.add(xb) D, I = index.search(xq, 10) self.assertEqual(D.shape, (10, 10)) self.assertEqual(I.shape, (10, 10)) def test_non_simd_dims(self): """Test dimensions not aligned to SIMD width (8/16).""" for d in [7, 9, 15, 17, 31, 33, 63, 65]: with self.subTest(d=d): ds = SyntheticDataset(d=d, nt=0, nb=100, nq=5, seed=42) xb = ds.get_database() xq = ds.get_queries() index = faiss.index_factory(d, "SQ8") index.train(xb) index.add(xb) D, I = index.search(xq, 10) self.assertEqual(D.shape, (5, 10)) self.assertEqual(I.shape, (5, 10)) def test_tqmse_non_simd_dims(self): factory_strings = [ "SQtqmse1", "SQtqmse2", "SQtqmse3", "SQtqmse4", "SQtqmse8", ] for d in [7, 9, 15, 17, 31, 33, 63, 65]: for factory_str in factory_strings: with self.subTest(d=d, factory_str=factory_str): ds = SyntheticDataset(d=d, nt=0, nb=100, nq=5, seed=42) xb = ds.get_database() xq = ds.get_queries() faiss.normalize_L2(xb) faiss.normalize_L2(xq) index = faiss.index_factory(d, factory_str) index.train(xb) index.add(xb) D, I = index.search(xq, 10) self.assertEqual(D.shape, (5, 10)) self.assertEqual(I.shape, (5, 10)) @for_all_simd_levels class TestScalarQuantizerIP(unittest.TestCase): def test_inner_product(self): d = 32 ds = SyntheticDataset(d=d, nt=0, nb=1000, nq=10, seed=42) xb = ds.get_database().copy() xq = ds.get_queries().copy() faiss.normalize_L2(xb) faiss.normalize_L2(xq) index = faiss.index_factory(d, "SQ8", faiss.METRIC_INNER_PRODUCT) index.train(xb) index.add(xb) D, _ = index.search(xq, 10) # normalized vectors: max IP should be close to 1 self.assertTrue(np.all(D[:, 0] <= 1.1)) self.assertTrue(np.all(D[:, 0] >= 0.5)) class TestTurboQUnbiasedness(unittest.TestCase): """Verify TurboQuant reconstruction preserves inner products.""" def test_reconstruction_quality(self): rng = np.random.RandomState(42) print("\n=== TurboQuant reconstruction quality ===") print(f"{'d':>6} {'qt':>8} {'self_ip':>10} {'cos_sim':>10}") print("-" * 42) for d in [128, 256, 384, 512, 768]: n = 500 xb = rng.randn(n, d).astype("float32") faiss.normalize_L2(xb) for qt_name, qt in [ ("SQtq2", faiss.ScalarQuantizer.QT_2bit_tq), ("SQtq3", faiss.ScalarQuantizer.QT_3bit_tq), ]: index = faiss.IndexScalarQuantizer(d, qt) index.train(xb) index.add(xb) xr = np.vstack([index.reconstruct(i) for i in range(n)]) # Self inner product: for each vector self_ip = np.array([xb[i] @ xr[i] for i in range(n)]) mean_self_ip = self_ip.mean() # Cosine similarity xr_norms = np.linalg.norm(xr, axis=1) cos_sim = np.array( [xb[i] @ xr[i] / (xr_norms[i] + 1e-30) for i in range(n)] ).mean() print( f"{d:>6} {qt_name:>8} " f"{mean_self_ip:>10.4f} {cos_sim:>10.4f}" ) # Reconstruction should point in the right direction self.assertGreater( mean_self_ip, 0.1, msg=f"d={d} {qt_name}: self IP {mean_self_ip:.4f} too low", ) class TestTurboQFullFactory(unittest.TestCase): """Test factory string construction for TurboQ full types.""" def test_factory_sqtq(self): cases = [ ("SQtq2", faiss.ScalarQuantizer.QT_2bit_tq), ("SQtq3", faiss.ScalarQuantizer.QT_3bit_tq), ("SQtq4", faiss.ScalarQuantizer.QT_4bit_tq), ("SQtq5", faiss.ScalarQuantizer.QT_5bit_tq), ] for factory_str, qtype in cases: with self.subTest(factory_str=factory_str): index = faiss.index_factory(32, factory_str) self.assertIsInstance(index, faiss.IndexScalarQuantizer) self.assertEqual(index.sq.qtype, qtype) def test_factory_ivf_sqtq(self): d = 64 for sq_str in ["SQtq2", "SQtq3", "SQtq4"]: with self.subTest(sq_str=sq_str): factory_str = f"IVF16,{sq_str}" index = faiss.index_factory(d, factory_str) self.assertIsInstance(index, faiss.IndexIVFScalarQuantizer) ivf = faiss.downcast_index(index) self.assertFalse(ivf.by_residual) class TestTurboQFullEncodeDecode(unittest.TestCase): """Encode/decode roundtrip for TurboQ full types.""" def test_encode_decode_roundtrip(self): for d in [32, 64, 128]: for qtype, max_err in [ (faiss.ScalarQuantizer.QT_2bit_tq, 1.5), (faiss.ScalarQuantizer.QT_3bit_tq, 1.2), (faiss.ScalarQuantizer.QT_4bit_tq, 0.8), ]: with self.subTest(d=d, qtype=qtype): rng = np.random.RandomState(42) xb = rng.randn(200, d).astype("float32") faiss.normalize_L2(xb) sq = faiss.ScalarQuantizer(d, qtype) sq.train(xb) codes = sq.compute_codes(xb) xr = sq.decode(codes) self.assertEqual(xr.shape, xb.shape) err = np.abs(xb - xr).max() self.assertLess(err, max_err) def test_non_simd_aligned_dims(self): for d in [7, 9, 15, 17, 33, 65]: with self.subTest(d=d): rng = np.random.RandomState(42) xb = rng.randn(100, d).astype("float32") faiss.normalize_L2(xb) sq = faiss.ScalarQuantizer(d, faiss.ScalarQuantizer.QT_2bit_tq) sq.train(xb) codes = sq.compute_codes(xb) xr = sq.decode(codes) self.assertEqual(xr.shape, xb.shape) class TestTurboQFullSearch(unittest.TestCase): """Search recall tests for TurboQ full types.""" def setUp(self): self.d = 64 self.rng = np.random.RandomState(42) self.nb = 5000 self.nq = 50 self.xb = self.rng.randn(self.nb, self.d).astype("float32") self.xq = self.rng.randn(self.nq, self.d).astype("float32") faiss.normalize_L2(self.xb) faiss.normalize_L2(self.xq) gt_index = faiss.IndexFlatL2(self.d) gt_index.add(self.xb) _, self.I_gt = gt_index.search(self.xq, 10) def _recall_at_1(self, index): _, I = index.search(self.xq, 10) return (self.I_gt[:, 0] == I[:, 0]).mean() def test_flat_sqtq_search(self): for sq_str in ["SQtq2", "SQtq3", "SQtq4"]: with self.subTest(sq_str=sq_str): index = faiss.index_factory(self.d, sq_str) index.train(self.xb) index.add(self.xb) recall = self._recall_at_1(index) self.assertGreater(recall, 0.05) def test_recall_increases_with_bits(self): recalls = {} for sq_str in ["SQtq2", "SQtq3", "SQtq4"]: index = faiss.index_factory(self.d, sq_str) index.train(self.xb) index.add(self.xb) recalls[sq_str] = self._recall_at_1(index) self.assertGreaterEqual(recalls["SQtq3"], recalls["SQtq2"] - 0.05) self.assertGreaterEqual(recalls["SQtq4"], recalls["SQtq3"] - 0.05) def test_ivf_sqtq_search(self): for sq_str in ["SQtq2", "SQtq3"]: with self.subTest(sq_str=sq_str): factory_str = f"IVF32,{sq_str}" index = faiss.index_factory(self.d, factory_str) index.train(self.xb) index.add(self.xb) index.nprobe = 8 _, I = index.search(self.xq, 10) self.assertEqual(I.shape, (self.nq, 10)) self.assertTrue(np.all(I[:, 0] >= 0)) class TestTurboQFullSerialization(unittest.TestCase): """Serialization roundtrip for TurboQ full types.""" def _roundtrip(self, index): fd, fname = tempfile.mkstemp() os.close(fd) try: faiss.write_index(index, fname) index2 = faiss.read_index(fname) return index2 finally: os.unlink(fname) def test_flat_serialize_roundtrip(self): d = 64 rng = np.random.RandomState(42) xb = rng.randn(500, d).astype("float32") xq = rng.randn(10, d).astype("float32") faiss.normalize_L2(xb) faiss.normalize_L2(xq) for qtype in [ faiss.ScalarQuantizer.QT_2bit_tq, faiss.ScalarQuantizer.QT_3bit_tq, faiss.ScalarQuantizer.QT_4bit_tq, faiss.ScalarQuantizer.QT_5bit_tq, ]: with self.subTest(qtype=qtype): index = faiss.IndexScalarQuantizer(d, qtype) index.train(xb) index.add(xb) D1, I1 = index.search(xq, 5) index2 = self._roundtrip(index) D2, I2 = index2.search(xq, 5) np.testing.assert_array_equal(I1, I2) np.testing.assert_allclose(D1, D2, atol=1e-6) def test_ivf_serialize_roundtrip(self): d = 64 rng = np.random.RandomState(42) xb = rng.randn(1000, d).astype("float32") xq = rng.randn(10, d).astype("float32") faiss.normalize_L2(xb) faiss.normalize_L2(xq) for sq_str in ["SQtq2", "SQtq3"]: with self.subTest(sq_str=sq_str): index = faiss.index_factory(d, f"IVF16,{sq_str}") index.train(xb) index.add(xb) index.nprobe = 4 D1, I1 = index.search(xq, 5) index2 = self._roundtrip(index) index2.nprobe = 4 D2, I2 = index2.search(xq, 5) np.testing.assert_array_equal(I1, I2) np.testing.assert_allclose(D1, D2, atol=1e-6) def test_serialize_preserves_qtype(self): d = 32 rng = np.random.RandomState(42) xb = rng.randn(100, d).astype("float32") faiss.normalize_L2(xb) for qtype in [ faiss.ScalarQuantizer.QT_2bit_tq, faiss.ScalarQuantizer.QT_3bit_tq, ]: with self.subTest(qtype=qtype): index = faiss.IndexScalarQuantizer(d, qtype) index.train(xb) index.add(xb) index2 = self._roundtrip(index) self.assertEqual(index2.sq.qtype, qtype) self.assertEqual(index2.ntotal, index.ntotal) self.assertEqual(index2.d, d) class TestTurboQFullDistances(unittest.TestCase): """Verify distance computation is consistent for TurboQ.""" def test_search_distances_are_finite(self): d = 64 rng = np.random.RandomState(42) xb = rng.randn(500, d).astype("float32") xq = rng.randn(10, d).astype("float32") faiss.normalize_L2(xb) faiss.normalize_L2(xq) for qtype in [ faiss.ScalarQuantizer.QT_2bit_tq, faiss.ScalarQuantizer.QT_3bit_tq, ]: with self.subTest(qtype=qtype): index = faiss.IndexScalarQuantizer(d, qtype) index.train(xb) index.add(xb) D, I = index.search(xq, 5) self.assertTrue(np.all(np.isfinite(D))) self.assertTrue(np.all(D[:, 0] >= 0)) self.assertTrue(np.all(I >= 0)) def test_search_distances_sorted(self): d = 64 rng = np.random.RandomState(42) xb = rng.randn(500, d).astype("float32") xq = rng.randn(10, d).astype("float32") faiss.normalize_L2(xb) faiss.normalize_L2(xq) for qtype in [ faiss.ScalarQuantizer.QT_2bit_tq, faiss.ScalarQuantizer.QT_3bit_tq, ]: with self.subTest(qtype=qtype): index = faiss.IndexScalarQuantizer(d, qtype) index.train(xb) index.add(xb) D, _ = index.search(xq, 10) for q in range(len(xq)): for k in range(1, 10): self.assertLessEqual(D[q, k - 1], D[q, k]) @for_all_simd_levels class TestSQByteDirectAcrossSIMDLevels(unittest.TestCase): """Every SIMD level must return the exact QT_8bit_direct{,_signed} distances. The distances are integers, so this class computes them from the definition and compares. An earlier version measured a second time at SIMDLevel.NONE and compared the two runs. A static build reports only the one level it compiles, and omits NONE, so that version skipped itself in the build shape that ships. d = 16 and d = 32 pass the d % 16 == 0 test in sq-dispatch.h. That test routes the query to DistanceComputerByte and DistanceComputerByteSigned. Two dispatch chains reach those kernels, and this class covers both. search() uses sq_select_InvertedListScanner. get_distance_computer() uses sq_select_distance_computer, which is the chain Refine(SQ8) follows. """ def reference(self, xq, xb, metric): """Returns the exact distances. The values are small integers. Every product and sum stays below 2**24, so the float32 result loses nothing.""" a = xq.astype("int64") b = xb.astype("int64") if metric == faiss.METRIC_L2: return ((a[:, None, :] - b[None, :, :]) ** 2).sum(-1) return a @ b.T def do_test(self, measure, expect): rng = np.random.RandomState(1234) for qtype in ( faiss.ScalarQuantizer.QT_8bit_direct, faiss.ScalarQuantizer.QT_8bit_direct_signed, ): lo, hi = ( (-128, 128) if qtype == faiss.ScalarQuantizer.QT_8bit_direct_signed else (0, 256) ) for metric in (faiss.METRIC_L2, faiss.METRIC_INNER_PRODUCT): for d in (16, 32): with self.subTest(qtype=qtype, metric=metric, d=d): xb = rng.randint(lo, hi, (500, d)).astype("float32") xq = rng.randint(lo, hi, (10, d)).astype("float32") index = faiss.IndexScalarQuantizer(d, qtype, metric) index.add(xb) got = measure(index, xq) truth = self.reference(xq, xb, metric) np.testing.assert_array_equal(got, expect(truth, metric)) def test_scanner(self): k = 10 def measure(index, xq): return index.search(xq, k)[0] def expect(truth, metric): # search returns L2 ascending and inner product descending. ordered = np.sort(truth, axis=1) if metric == faiss.METRIC_INNER_PRODUCT: ordered = ordered[:, ::-1] return ordered[:, :k] self.do_test(measure, expect) def test_distance_computer(self): def measure(index, xq): dc = index.get_distance_computer() out = np.empty((len(xq), index.ntotal), dtype="float32") for q in range(len(xq)): dc.set_query(faiss.swig_ptr(xq[q])) for i in range(index.ntotal): out[q, i] = dc(int(i)) return out self.do_test(measure, lambda truth, metric: truth)