usearch 0.13.0 → 0.15.0

This diff represents the content of publicly available package versions that have been released to one of the supported registries. The information contained in this diff is provided for informational purposes only and reflects changes between package versions as they appear in their respective public registries.
Files changed (64) hide show
  1. package/.github/workflows/release.yml +62 -7
  2. package/.github/workflows/update_version.sh +6 -1
  3. package/.vscode/launch.json +21 -22
  4. package/.vscode/settings.json +7 -1
  5. package/CITATION.cff +4 -4
  6. package/CMakeLists.txt +1 -1
  7. package/Cargo.toml +1 -1
  8. package/Package.swift +0 -1
  9. package/README.md +100 -27
  10. package/binding.gyp +0 -1
  11. package/build.gradle +1 -1
  12. package/build.rs +3 -3
  13. package/c/lib.cpp +32 -9
  14. package/c/usearch.h +67 -36
  15. package/conanfile.py +1 -1
  16. package/{src → cpp}/bench.cpp +52 -36
  17. package/{src → cpp}/test.cpp +17 -7
  18. package/docs/benchmarks.md +33 -10
  19. package/docs/c/index.rst +11 -0
  20. package/docs/c/reference.rst +7 -0
  21. package/docs/compilation.md +24 -8
  22. package/docs/conf.dox +6 -6
  23. package/docs/conf.py +3 -1
  24. package/docs/cpp/index.rst +11 -0
  25. package/docs/cpp/reference.rst +7 -0
  26. package/docs/index.rst +10 -0
  27. package/docs/java/index.rst +11 -0
  28. package/docs/java/reference.rst +7 -0
  29. package/docs/javascript/index.rst +11 -0
  30. package/docs/javascript/reference.rst +9 -0
  31. package/docs/python/index.rst +11 -0
  32. package/docs/python/reference.rst +35 -0
  33. package/golang/lib.go +1 -1
  34. package/include/usearch/{usearch.hpp → index.hpp} +152 -148
  35. package/include/usearch/index_punned_dense.hpp +633 -0
  36. package/include/usearch/index_punned_helpers.hpp +528 -0
  37. package/java/cloud/unum/usearch/cloud_unum_usearch_Index.cpp +8 -7
  38. package/javascript/lib.cpp +26 -50
  39. package/javascript/usearch.js +117 -2
  40. package/objc/USearchObjective.mm +145 -49
  41. package/objc/include/USearchObjective.h +106 -30
  42. package/package.json +2 -2
  43. package/python/lib.cpp +178 -320
  44. package/python/scripts/bench.ipynb +170 -0
  45. package/python/scripts/bench.py +102 -185
  46. package/python/scripts/index_faiss.py +100 -0
  47. package/python/scripts/test.py +57 -34
  48. package/python/usearch/eval.py +353 -0
  49. package/python/usearch/index.py +138 -117
  50. package/python/usearch/io.py +7 -2
  51. package/python/usearch/join.py +100 -0
  52. package/python/usearch/numba.py +101 -0
  53. package/rust/lib.cpp +27 -14
  54. package/rust/lib.hpp +4 -10
  55. package/rust/lib.rs +50 -17
  56. package/setup.py +12 -2
  57. package/simsimd/include/simsimd/simsimd.h +32 -3
  58. package/swift/Index+Sugar.swift +55 -29
  59. package/swift/Test.swift +1 -1
  60. package/wolfram/lib.cpp +4 -4
  61. package/python/scripts/bench_params.py +0 -161
  62. package/python/usearch/synthetic.py +0 -32
  63. package/src/punned.hpp +0 -887
  64. /package/{src → cpp}/CMakeLists.txt +0 -0
@@ -4,28 +4,60 @@
4
4
  # into the primary `Index` class, connecting USearch with Numba.
5
5
  import os
6
6
  from math import sqrt
7
- from typing import Optional, Callable, Union, NamedTuple, List
7
+ from typing import Optional, Callable, Union, NamedTuple, List, Iterable
8
8
 
9
9
  import numpy as np
10
10
 
11
11
  from usearch.compiled import Index as _CompiledIndex
12
- from usearch.compiled import SetsIndex as _CompiledSetsIndex
13
- from usearch.compiled import BitsIndex as _CompiledBitsIndex
12
+ from usearch.compiled import SparseIndex as _CompiledSetsIndex
14
13
 
15
- from usearch.compiled import MetricKind
14
+ from usearch.compiled import MetricKind, ScalarKind
16
15
  from usearch.compiled import (
17
16
  DEFAULT_CONNECTIVITY,
18
17
  DEFAULT_EXPANSION_ADD,
19
18
  DEFAULT_EXPANSION_SEARCH,
19
+
20
+ USES_OPENMP,
21
+ USES_SIMSIMD,
20
22
  )
21
23
 
22
- BitwiseMetricKind = (
23
- MetricKind.BitwiseHamming,
24
- MetricKind.BitwiseTanimoto,
25
- MetricKind.BitwiseSorensen,
24
+ MetricKindBitwise = (
25
+ MetricKind.Hamming,
26
+ MetricKind.Tanimoto,
27
+ MetricKind.Sorensen,
26
28
  )
27
29
 
28
- SetsIndex = _CompiledSetsIndex
30
+ SparseIndex = _CompiledSetsIndex
31
+
32
+ Label = np.longlong
33
+
34
+
35
+ def _normalize_dtype(dtype) -> ScalarKind:
36
+ if dtype is None:
37
+ return ScalarKind.F32
38
+ if isinstance(dtype, ScalarKind):
39
+ return dtype
40
+
41
+ if isinstance(dtype, str):
42
+ _normalize = {
43
+ 'f64': ScalarKind.F64,
44
+ 'f32': ScalarKind.F32,
45
+ 'f16': ScalarKind.F16,
46
+ 'f8': ScalarKind.F8,
47
+ 'b1': ScalarKind.B1,
48
+ }
49
+ return _normalize[dtype.lower()]
50
+
51
+ if not isinstance(dtype, ScalarKind):
52
+ _normalize = {
53
+ np.float64: ScalarKind.F64,
54
+ np.float32: ScalarKind.F32,
55
+ np.float16: ScalarKind.F16,
56
+ np.int8: ScalarKind.F8,
57
+ }
58
+ return _normalize[dtype]
59
+
60
+ return dtype
29
61
 
30
62
 
31
63
  class Matches(NamedTuple):
@@ -65,83 +97,15 @@ class Matches(NamedTuple):
65
97
  for label, distance in zip(labels, distances)
66
98
  ]
67
99
 
100
+ def recall_first(self, expected: np.ndarray) -> float:
101
+ best_matches = self.labels if not self.is_batch else self.labels[:, 0]
102
+ return np.sum(best_matches == expected) / len(expected)
103
+
68
104
  def __repr__(self) -> str:
69
105
  return f'usearch.Matches({self.total_matches})' if self.is_batch else \
70
106
  f'usearch.Matches({self.total_matches} across {self.batch_size} queries)'
71
107
 
72
108
 
73
- def jit_metric(ndim: int, metric: MetricKind, dtype: str = 'f32') -> Callable:
74
-
75
- try:
76
- from numba import cfunc, types, carray
77
- except ImportError:
78
- raise ModuleNotFoundError(
79
- 'To use JIT install Numba with `pip install numba`.'
80
- 'Alternatively, reinstall usearch with `pip install usearch[jit]`')
81
-
82
- # Showcases how to use Numba to JIT-compile similarity measures for USearch.
83
- # https://numba.readthedocs.io/en/stable/reference/jit-compilation.html#c-callbacks
84
-
85
- if dtype == 'f32':
86
- signature = types.float32(
87
- types.CPointer(types.float32),
88
- types.CPointer(types.float32),
89
- types.uint64, types.uint64)
90
-
91
- if metric == MetricKind.IP:
92
-
93
- @cfunc(signature)
94
- def numba_ip(a, b, _n, _m):
95
- a_array = carray(a, ndim)
96
- b_array = carray(b, ndim)
97
- ab = 0.0
98
- for i in range(ndim):
99
- ab += a_array[i] * b_array[i]
100
- return 1 - ab
101
-
102
- return numba_ip
103
-
104
- if metric == MetricKind.Cos:
105
-
106
- @cfunc(signature)
107
- def numba_cos(a, b, _n, _m):
108
- a_array = carray(a, ndim)
109
- b_array = carray(b, ndim)
110
- ab = 0.0
111
- a_sq = 0.0
112
- b_sq = 0.0
113
- for i in range(ndim):
114
- ab += a_array[i] * b_array[i]
115
- a_sq += a_array[i] * a_array[i]
116
- b_sq += b_array[i] * b_array[i]
117
- a_norm = sqrt(a_sq)
118
- b_norm = sqrt(b_sq)
119
- if a_norm == 0 and b_norm == 0:
120
- return 0
121
- elif a_norm == 0 or b_norm == 0 or ab == 0:
122
- return 1
123
- else:
124
- return 1 - ab / (a_norm * b_norm)
125
-
126
- return numba_cos
127
-
128
- if metric == MetricKind.L2sq:
129
-
130
- @cfunc(signature)
131
- def numba_l2sq(a, b, _n, _m):
132
- a_array = carray(a, ndim)
133
- b_array = carray(b, ndim)
134
- ab_delta_sq = types.float32(0.0)
135
- for i in range(ndim):
136
- ab_delta_sq += (a_array[i] - b_array[i]) * \
137
- (a_array[i] - b_array[i])
138
- return ab_delta_sq
139
-
140
- return numba_l2sq
141
-
142
- return None
143
-
144
-
145
109
  class Index:
146
110
  """Fast JIT-compiled vector-search index for dense equi-dimensional embeddings.
147
111
 
@@ -155,7 +119,7 @@ class Index:
155
119
  def __init__(
156
120
  self,
157
121
  ndim: int,
158
- metric: Union[MetricKind, Callable] = MetricKind.IP,
122
+ metric: Union[MetricKind, Callable, str] = MetricKind.IP,
159
123
  dtype: Optional[str] = None,
160
124
  jit: bool = False,
161
125
 
@@ -173,11 +137,11 @@ class Index:
173
137
  :type ndim: int
174
138
 
175
139
  :param metric: Distance function, defaults to MetricKind.IP
176
- :type metric: Union[MetricKind, Callable], optional
140
+ :type metric: Union[MetricKind, Callable, str], optional
177
141
  Kind of the distance function, or the Numba `cfunc` JIT-compiled object.
178
142
  Possible `MetricKind` values: IP, Cos, L2sq, Haversine, Pearson,
179
- BitwiseHamming, BitwiseTanimoto, BitwiseSorensen.
180
- Not every kind is JIT-able. For Jaccard distance, use `SetsIndex`.
143
+ Hamming, Tanimoto, Sorensen.
144
+ Not every kind is JIT-able. For Jaccard distance, use `SparseIndex`.
181
145
 
182
146
  :param dtype: Scalar type for internal vector storage, defaults to None
183
147
  :type dtype: str, optional
@@ -221,6 +185,18 @@ class Index:
221
185
 
222
186
  if metric is None:
223
187
  metric = MetricKind.IP
188
+ elif isinstance(metric, str):
189
+ _normalize = {
190
+ 'cos': MetricKind.Cos,
191
+ 'ip': MetricKind.IP,
192
+ 'l2_sq': MetricKind.L2sq,
193
+ 'haversine': MetricKind.Haversine,
194
+ 'perason': MetricKind.Pearson,
195
+ 'hamming': MetricKind.Hamming,
196
+ 'tanimoto': MetricKind.Tanimoto,
197
+ 'sorensen': MetricKind.Sorensen,
198
+ }
199
+ metric = _normalize[metric.lower()]
224
200
 
225
201
  if isinstance(metric, Callable):
226
202
  self._metric_kind = MetricKind.Unknown
@@ -229,10 +205,18 @@ class Index:
229
205
 
230
206
  elif isinstance(metric, MetricKind):
231
207
  if jit:
208
+
209
+ try:
210
+ from usearch.numba import jit
211
+ except ImportError:
212
+ raise ModuleNotFoundError(
213
+ 'To use JIT install Numba with `pip install numba`.'
214
+ 'Alternatively, reinstall usearch with `pip install usearch[jit]`')
215
+
232
216
  self._metric_kind = metric
233
- self._metric_jit = jit_metric(
217
+ self._metric_jit = jit(
234
218
  ndim=ndim,
235
- metric_kind=metric,
219
+ metric=metric,
236
220
  dtype=dtype,
237
221
  )
238
222
  self._metric_pointer = self._metric_jit.address if \
@@ -245,27 +229,22 @@ class Index:
245
229
  raise ValueError(
246
230
  'The `metric` must be Numba callback or a `MetricKind`')
247
231
 
248
- if metric in BitwiseMetricKind:
249
- self._compiled = _CompiledBitsIndex(
250
- bits=ndim,
251
- metric=self._metric_kind,
252
- connectivity=connectivity,
253
- expansion_add=expansion_add,
254
- expansion_search=expansion_search,
255
- )
232
+ if metric in MetricKindBitwise:
233
+ assert dtype is None or dtype == ScalarKind.B1
234
+ dtype = ScalarKind.B1
256
235
  else:
257
- if dtype is None:
258
- dtype = 'f32'
259
- self._compiled = _CompiledIndex(
260
- ndim=ndim,
261
- metric=self._metric_kind,
262
- metric_pointer=self._metric_pointer,
263
- dtype=dtype,
264
- connectivity=connectivity,
265
- expansion_add=expansion_add,
266
- expansion_search=expansion_search,
267
- tune=tune,
268
- )
236
+ dtype = _normalize_dtype(dtype)
237
+
238
+ self._compiled = _CompiledIndex(
239
+ ndim=ndim,
240
+ metric=self._metric_kind,
241
+ metric_pointer=self._metric_pointer,
242
+ dtype=dtype,
243
+ connectivity=connectivity,
244
+ expansion_add=expansion_add,
245
+ expansion_search=expansion_search,
246
+ tune=tune,
247
+ )
269
248
 
270
249
  self.path = path
271
250
  if path and os.path.exists(path):
@@ -276,7 +255,7 @@ class Index:
276
255
 
277
256
  def add(
278
257
  self, labels, vectors, *,
279
- copy: bool = True, threads: int = 0):
258
+ copy: bool = True, threads: int = 0) -> Union[int, np.ndarray]:
280
259
  """Inserts one or move vectors into the index.
281
260
 
282
261
  For maximal performance the `labels` and `vectors`
@@ -291,7 +270,7 @@ class Index:
291
270
  pass `copy=False`, if you can guarantee the lifetime of the
292
271
  primary vectors store during the process of construction.
293
272
 
294
- :param labels: Unique identifier for passed vectors.
273
+ :param labels: Unique identifier for passed vectors, optional
295
274
  :type labels: Buffer
296
275
  :param vectors: Collection of vectors.
297
276
  :type vectors: Buffer
@@ -299,6 +278,9 @@ class Index:
299
278
  :type copy: bool, optional
300
279
  :param threads: Optimal number of cores to use, defaults to 0
301
280
  :type threads: int, optional
281
+
282
+ :return: Inserted label or labels
283
+ :type: Union[int, np.ndarray]
302
284
  """
303
285
  assert isinstance(vectors, np.ndarray), 'Expects a NumPy array'
304
286
  assert vectors.ndim == 1 or vectors.ndim == 2, 'Expects a matrix or vector'
@@ -314,11 +296,12 @@ class Index:
314
296
  labels = start_id
315
297
  if isinstance(labels, np.ndarray):
316
298
  labels = labels.astype(np.longlong)
299
+ elif isinstance(labels, Iterable):
300
+ labels = np.array(labels, dtype=np.longlong)
317
301
 
318
302
  self._compiled.add(labels, vectors, copy=copy, threads=threads)
319
303
 
320
- if generate_labels:
321
- return labels
304
+ return labels
322
305
 
323
306
  def search(
324
307
  self, vectors, k: int = 10, *,
@@ -343,8 +326,24 @@ class Index:
343
326
  )
344
327
  return Matches(*tuple_)
345
328
 
329
+ @property
330
+ def specs(self) -> dict:
331
+ return {
332
+ 'Class': 'usearch.Index',
333
+ 'Connectivity': self.connectivity,
334
+ 'Size': self.size,
335
+ 'Dimensions': self.ndim,
336
+ 'Expansion@Add': self.expansion_add,
337
+ 'Expansion@Search': self.expansion_search,
338
+ 'OpenMP': USES_OPENMP,
339
+ 'SimSIMD': USES_SIMSIMD,
340
+ 'JIT': self.jit,
341
+ 'DType': self.dtype,
342
+ 'Path': self.path,
343
+ }
344
+
346
345
  def __len__(self) -> int:
347
- return len(self._compiled)
346
+ return self._compiled.__len__()
348
347
 
349
348
  @property
350
349
  def jit(self) -> bool:
@@ -363,7 +362,7 @@ class Index:
363
362
  return self._metric_kind
364
363
 
365
364
  @property
366
- def dtype(self) -> str:
365
+ def dtype(self) -> ScalarKind:
367
366
  return self._compiled.dtype
368
367
 
369
368
  @property
@@ -387,11 +386,11 @@ class Index:
387
386
  return self._compiled.expansion_search
388
387
 
389
388
  @expansion_add.setter
390
- def change_expansion_add(self, v: int):
389
+ def expansion_add(self, v: int):
391
390
  self._compiled.expansion_add = v
392
391
 
393
392
  @expansion_search.setter
394
- def change_expansion_search(self, v: int):
393
+ def expansion_search(self, v: int):
395
394
  self._compiled.expansion_search = v
396
395
 
397
396
  def save(self, path: os.PathLike):
@@ -406,5 +405,27 @@ class Index:
406
405
  def clear(self):
407
406
  self._compiled.clear()
408
407
 
409
- def remove(self, label: int):
410
- pass
408
+ @property
409
+ def labels(self) -> np.ndarray:
410
+ return self._compiled.labels
411
+
412
+ def get_vectors(
413
+ self,
414
+ labels: np.ndarray,
415
+ dtype: ScalarKind = ScalarKind.F32) -> np.ndarray:
416
+
417
+ dtype = _normalize_dtype(dtype)
418
+ return np.vstack([self._compiled.__getitem__(l, dtype) for l in labels])
419
+
420
+ @property
421
+ def vectors(self) -> np.ndarray:
422
+ return self.get_vectors(self.labels)
423
+
424
+ def __delitem__(self, label: int):
425
+ raise NotImplementedError()
426
+
427
+ def __contains__(self, label: int) -> bool:
428
+ return self._compiled.__contains__(label)
429
+
430
+ def __getitem__(self, label: int) -> np.ndarray:
431
+ return self._compiled.__getitem__(label)
@@ -7,7 +7,7 @@ def load_matrix(
7
7
  filename: str,
8
8
  start_row: int = 0, count_rows: int = None,
9
9
  view: bool = False) -> np.ndarray:
10
- """Read *.ibin, *.hbin, *.fbin, *.dbin files with matrixes.
10
+ """Read *.ibin, *.bbib, *.hbin, *.fbin, *.dbin files with matrices.
11
11
 
12
12
  :param filename: path to the matrix file
13
13
  :param start_row: start reading vectors from this index
@@ -31,6 +31,9 @@ def load_matrix(
31
31
  elif filename.endswith('.ibin'):
32
32
  dtype = np.int32
33
33
  scalar_size = 4
34
+ elif filename.endswith('.bbin'):
35
+ dtype = np.int8
36
+ scalar_size = 1
34
37
  else:
35
38
  raise Exception('Unknown file type')
36
39
 
@@ -49,7 +52,7 @@ def load_matrix(
49
52
 
50
53
 
51
54
  def save_matrix(vectors: np.ndarray, filename: str):
52
- """Write *.ibin, *.hbin, *.fbin, *.dbin files with matrixes.
55
+ """Write *.ibin, *.bbib, *.hbin, *.fbin, *.dbin files with matrices.
53
56
 
54
57
  :param vectors: the matrix to serialize
55
58
  :type vectors: numpy.ndarray
@@ -65,6 +68,8 @@ def save_matrix(vectors: np.ndarray, filename: str):
65
68
  dtype = np.float16
66
69
  elif filename.endswith('.ibin'):
67
70
  dtype = np.int32
71
+ elif filename.endswith('.bbin'):
72
+ dtype = np.int8
68
73
  else:
69
74
  raise Exception('Unknown file type')
70
75
 
@@ -0,0 +1,100 @@
1
+ from collections import deque
2
+
3
+ import numpy as np
4
+ from numba import njit
5
+
6
+ from index import Index, Matches, Label
7
+
8
+
9
+ class bidict(dict):
10
+
11
+ def __init__(self, *args, **kwargs):
12
+ super(bidict, self).__init__(*args, **kwargs)
13
+ self.inverse = {}
14
+ for key, value in self.items():
15
+ self.inverse.setdefault(value, []).append(key)
16
+
17
+ def __setitem__(self, key, value):
18
+ if key in self:
19
+ self.inverse[self[key]].remove(key)
20
+ super(bidict, self).__setitem__(key, value)
21
+ self.inverse.setdefault(value, []).append(key)
22
+
23
+ def __delitem__(self, key):
24
+ self.inverse.setdefault(self[key], []).remove(key)
25
+ if self[key] in self.inverse and not self.inverse[self[key]]:
26
+ del self.inverse[self[key]]
27
+ super(bidict, self).__delitem__(key)
28
+
29
+
30
+ @njit
31
+ def index_in_array(array: np.ndarray, item: Label) -> int:
32
+ for idx, val in np.ndenumerate(array):
33
+ if val == item:
34
+ return idx
35
+
36
+ return len(array)
37
+
38
+
39
+ def semantic_join(a: Index, b: Index, resolution: int = 10) -> bidict:
40
+ """Performs a Semantic Join, matching one entry from `a` with
41
+ one entry from `b`, converging towards a Stable Marriage.
42
+ Assuming the collections can be different in size, classical solution
43
+ doesn't provide
44
+
45
+ :param resolution: Approximate matches per member to consider, defaults to 10
46
+ :type resolution: int, optional
47
+ :return: Bidirectional mapping from men to women and `.inverse`
48
+ :rtype: bidict
49
+ """
50
+
51
+ men, women = (a, b) if len(a) > len(b) else (b, a)
52
+ matches = bidict()
53
+
54
+ men_vectors = np.vstack(men[i] for i in range(len(men)))
55
+ women_vectors = np.vstack(women[i] for i in range(len(women)))
56
+
57
+ man_count_proposed = np.zeros(len(men))
58
+ man_to_women_preferences: Matches = men.search(women_vectors, resolution)
59
+ woman_to_men_preferences: Matches = women.search(men_vectors, resolution)
60
+
61
+ # A nice optimization is to prioritize free men by the quality of their
62
+ # best remaining variant. TODO: Replace with `heapq`
63
+ free_men = deque(range(len(men)))
64
+
65
+ idle_cycles: int = 0
66
+ hopeless_men_count = len(men) - len(women)
67
+ while len(free_men) > hopeless_men_count:
68
+
69
+ # In the worst case scenario we may need to match more candidates for every
70
+ # remaining man. This, however, may drastically increase the runtime.
71
+ if len(free_men) == idle_cycles:
72
+ break
73
+
74
+ man: Label = free_men.popleft()
75
+ count_proposals = man_count_proposed[man]
76
+ if count_proposals == man_to_women_preferences.counts[man]:
77
+ free_men.append(man)
78
+ idle_cycles += 1
79
+ continue
80
+
81
+ woman: Label = man_to_women_preferences.labels[man, count_proposals]
82
+ man_count_proposed[man] += 1
83
+
84
+ if woman not in matches.inverse:
85
+ matches[man] = woman
86
+
87
+ else:
88
+ her_preferences = woman_to_men_preferences.labels[woman, :]
89
+ husband: Label = matches.inverse[woman]
90
+ husband_rank = index_in_array(her_preferences, husband)
91
+ challenger_rank = index_in_array(her_preferences, man)
92
+ if challenger_rank < husband_rank:
93
+ del matches[husband]
94
+ free_men.append(husband)
95
+ matches[man] = woman
96
+
97
+ else:
98
+ free_men.append(man)
99
+
100
+ return matches
@@ -0,0 +1,101 @@
1
+ # The purpose of this file is to provide Pythonic wrapper on top
2
+ # the native precompiled CPython module. It improves compatibility
3
+ # Python tooling, linters, and static analyzers. It also embeds JIT
4
+ # into the primary `Index` class, connecting USearch with Numba.
5
+ from math import sqrt
6
+ from typing import Callable
7
+
8
+ import numpy as np
9
+ from numba import cfunc, types, carray
10
+
11
+ from usearch.compiled import MetricKind
12
+
13
+
14
+ signature_f16args = types.float32(types.CPointer(
15
+ types.float16), types.CPointer(types.float16))
16
+
17
+ signature_f32args = types.float32(types.CPointer(
18
+ types.float32), types.CPointer(types.float32))
19
+
20
+ signature_f64args = types.float32(types.CPointer(
21
+ types.float64), types.CPointer(types.float64))
22
+
23
+
24
+ def jit(ndim: int,
25
+ metric: MetricKind = MetricKind.Cos,
26
+ dtype: np.generic = np.float32) -> Callable:
27
+ """JIT-compiles a distance metric specifically tuned for the target hardware
28
+ and number of dimensions.
29
+
30
+ Uses Numba `cfunc` functionality, annotating it with Numba `types` instead
31
+ of `ctypes` to support half-precision.
32
+ https://numba.readthedocs.io/en/stable/reference/jit-compilation.html#c-callbacks
33
+ """
34
+ normalize_dtype = {
35
+ # Half-precision is still unsupported
36
+ # https://github.com/numba/numba/issues/4402
37
+ # 'f16': np.float16,
38
+ 'f32': np.float32,
39
+ 'f64': np.float64,
40
+ }
41
+ if dtype in normalize_dtype.keys():
42
+ dtype = normalize_dtype[dtype]
43
+ if dtype not in normalize_dtype.values():
44
+ return None
45
+
46
+ scalar_kind_to_accumulator_type = {
47
+ np.float16: types.float32,
48
+ np.float32: types.float32,
49
+ np.float64: types.float64,
50
+ }
51
+ accumulator = scalar_kind_to_accumulator_type[dtype]
52
+
53
+ def numba_ip(a, b):
54
+ a_array = carray(a, ndim)
55
+ b_array = carray(b, ndim)
56
+ ab = accumulator(0)
57
+ for i in range(ndim):
58
+ ab += a_array[i] * b_array[i]
59
+ return types.float32(1 - ab)
60
+
61
+ def numba_cos(a, b):
62
+ a_array = carray(a, ndim)
63
+ b_array = carray(b, ndim)
64
+ ab = accumulator(0)
65
+ a_sq = accumulator(0)
66
+ b_sq = accumulator(0)
67
+ for i in range(ndim):
68
+ ab += a_array[i] * b_array[i]
69
+ a_sq += a_array[i] * a_array[i]
70
+ b_sq += b_array[i] * b_array[i]
71
+ a_norm = sqrt(a_sq)
72
+ b_norm = sqrt(b_sq)
73
+ if a_norm == 0 and b_norm == 0:
74
+ return types.float32(0)
75
+ elif a_norm == 0 or b_norm == 0 or ab == 0:
76
+ return types.float32(1)
77
+ else:
78
+ return types.float32(1 - ab / (a_norm * b_norm))
79
+
80
+ def numba_l2sq(a, b):
81
+ a_array = carray(a, ndim)
82
+ b_array = carray(b, ndim)
83
+ ab_delta_sq = accumulator(0)
84
+ for i in range(ndim):
85
+ ab_delta_sq += (a_array[i] - b_array[i]) * \
86
+ (a_array[i] - b_array[i])
87
+ return types.float32(ab_delta_sq)
88
+
89
+ scalar_kind_to_signature = {
90
+ np.float16: signature_f16args,
91
+ np.float32: signature_f32args,
92
+ np.float64: signature_f64args,
93
+ }
94
+
95
+ metric_kind_to_function = {
96
+ MetricKind.IP: numba_ip,
97
+ MetricKind.Cos: numba_cos,
98
+ MetricKind.L2sq: numba_l2sq,
99
+ }
100
+
101
+ return cfunc(scalar_kind_to_signature[dtype])(metric_kind_to_function[metric])
package/rust/lib.cpp CHANGED
@@ -62,7 +62,7 @@ void Index::save(rust::Str path) const { index_->save(std::string(path).c_str())
62
62
  void Index::load(rust::Str path) const { index_->load(std::string(path).c_str()); }
63
63
  void Index::view(rust::Str path) const { index_->view(std::string(path).c_str()); }
64
64
 
65
- accuracy_t accuracy(rust::Str quant) { return accuracy_from_name(quant.data(), quant.size()); }
65
+ scalar_kind_t accuracy(rust::Str quant) { return scalar_kind_from_name(quant.data(), quant.size()); }
66
66
 
67
67
  std::unique_ptr<Index> wrap(punned_t&& index) {
68
68
  std::unique_ptr<punned_t> punned_ptr;
@@ -72,26 +72,39 @@ std::unique_ptr<Index> wrap(punned_t&& index) {
72
72
  return result;
73
73
  }
74
74
 
75
- index_config_t config(size_t connectivity, size_t exp_add, size_t exp_search) {
75
+ index_config_t config(size_t connectivity) {
76
76
  index_config_t result;
77
77
  result.connectivity = connectivity ? connectivity : default_connectivity();
78
- result.expansion_add = exp_add ? exp_add : default_expansion_add();
79
- result.expansion_search = exp_search ? exp_search : default_expansion_search();
80
78
  return result;
81
79
  }
82
80
 
83
- std::unique_ptr<Index> new_ip(size_t dims, rust::Str quant, size_t connectivity, size_t exp_add, size_t exp_search) {
84
- return wrap(punned_t::ip(dims, accuracy(quant), config(connectivity, exp_add, exp_search)));
81
+ metric_kind_t rust_to_cpp_metric(MetricKind value) {
82
+ switch (value) {
83
+ case MetricKind::IP: return metric_kind_t::ip_k;
84
+ case MetricKind::L2Sq: return metric_kind_t::l2sq_k;
85
+ case MetricKind::Cos: return metric_kind_t::cos_k;
86
+ case MetricKind::Pearson: return metric_kind_t::pearson_k;
87
+ case MetricKind::Haversine: return metric_kind_t::haversine_k;
88
+ case MetricKind::Hamming: return metric_kind_t::hamming_k;
89
+ case MetricKind::Tanimoto: return metric_kind_t::tanimoto_k;
90
+ case MetricKind::Sorensen: return metric_kind_t::sorensen_k;
91
+ default: return metric_kind_t::unknown_k;
92
+ }
85
93
  }
86
94
 
87
- std::unique_ptr<Index> new_l2sq(size_t dims, rust::Str quant, size_t connectivity, size_t exp_add, size_t exp_search) {
88
- return wrap(punned_t::l2sq(dims, accuracy(quant), config(connectivity, exp_add, exp_search)));
95
+ scalar_kind_t rust_to_cpp_scalar(ScalarKind value) {
96
+ switch (value) {
97
+ case ScalarKind::F8: return scalar_kind_t::f8_k;
98
+ case ScalarKind::F16: return scalar_kind_t::f16_k;
99
+ case ScalarKind::F32: return scalar_kind_t::f32_k;
100
+ case ScalarKind::F64: return scalar_kind_t::f64_k;
101
+ case ScalarKind::B1: return scalar_kind_t::b1x8_k;
102
+ default: return scalar_kind_t::unknown_k;
103
+ }
89
104
  }
90
105
 
91
- std::unique_ptr<Index> new_cos(size_t dims, rust::Str quant, size_t connectivity, size_t exp_add, size_t exp_search) {
92
- return wrap(punned_t::cos(dims, accuracy(quant), config(connectivity, exp_add, exp_search)));
93
- }
94
-
95
- std::unique_ptr<Index> new_haversine(rust::Str quant, size_t connectivity, size_t exp_add, size_t exp_search) {
96
- return wrap(punned_t::haversine(accuracy(quant), config(connectivity, exp_add, exp_search)));
106
+ std::unique_ptr<Index> new_index(IndexOptions const& options) {
107
+ return wrap(punned_t::make(options.dimensions, rust_to_cpp_metric(options.metric), config(options.connectivity),
108
+ rust_to_cpp_scalar(options.quantization), options.expansion_add,
109
+ options.expansion_search));
97
110
  }