usearch 0.19.1 → 0.19.2

package/Cargo.toml

@@ -1,6 +1,6 @@
 [package]
 name = "usearch"
-version = "0.19.1"
+version = "0.19.2"
 authors = ["Ashot Vardanian <1983160+ashvardanian@users.noreply.github.com>"]
 description = "Smaller & Faster Single-File Vector Search Engine from Unum"
 edition = "2021"

package/conanfile.py

@@ -4,7 +4,7 @@ from conans import ConanFile
 class USearchConan(ConanFile):
 
     name = 'USearch'
-    version = '0.19.1'
+    version = '0.19.2'
     license = 'Apache 2.0'
     url = 'https://github.com/unum-cloud/usearch'
     description = 'Smaller & Faster Single-File Vector Search Engine from Unum'

package/package.json

@@ -1,6 +1,6 @@
 {
     "name": "usearch",
-    "version": "0.19.1",
+    "version": "0.19.2",
     "description": "Smaller & Faster Single-File Vector Search Engine from Unum",
     "author": "Ashot Vardanian",
     "license": "Apache 2.0",

package/python/scripts/bench.py

@@ -1,14 +1,13 @@
-import os
-import logging
 import itertools
-from typing import List, Optional, Iterable
+from typing import List
 from dataclasses import asdict
 
 import numpy as np
 import pandas as pd
 
-from usearch.index import Index, Label
-from usearch.eval import Evaluation, Dataset, AddTask
+from usearch.index import Index, Label, MetricKind, ScalarKind
+from usearch.numba import jit as njit
+from usearch.eval import Evaluation, AddTask
 from usearch.index import (
     DEFAULT_CONNECTIVITY,
     DEFAULT_EXPANSION_ADD,
@@ -26,14 +25,15 @@ def bench_speed(
 ) -> pd.DataFrame:
     # Build various indexes:
     indexes = []
-    jit_options = [False]
-    if jit:
-        jit_options.append(jit)
-    dtype_options = ["f32", "f16", "f8"]
+    jit_options = [False, True] if jit else [False]
+    dtype_options = [ScalarKind.F32, ScalarKind.F16, ScalarKind.F8]
     for jit, dtype in itertools.product(jit_options, dtype_options):
+        metric = MetricKind.IP
+        if jit:
+            metric = njit(eval.ndim, metric, dtype)
         index = Index(
             ndim=eval.ndim,
-            jit=jit,
+            metric=metric,
             dtype=dtype,
             expansion_add=expansion_add,
             expansion_search=expansion_search,

package/python/scripts/test.py

@@ -1,9 +1,12 @@
 import os
+import importlib.util
+
 import pytest
 import numpy as np
 
 from usearch.io import load_matrix, save_matrix
-from usearch.eval import recall_members, random_vectors
+from usearch.eval import random_vectors
+from usearch.numba import jit as njit
 
 from usearch.index import (
     Index,
@@ -13,9 +16,7 @@ from usearch.index import (
     MetricSignature,
     CompiledMetric,
     Matches,
-    MetricKindBitwise,
 )
-from usearch.index import _normalize_dtype
 from usearch.index import (
     DEFAULT_CONNECTIVITY,
     DEFAULT_EXPANSION_ADD,
@@ -34,7 +35,7 @@ index_types = [
 numpy_types = [np.float32, np.float64, np.float16]
 
 connectivity_options = [3, 13, 50, DEFAULT_CONNECTIVITY]
-jit_options = [False]
+jit_options = [False, True] if importlib.util.find_spec("spam") is not None else [False]
 continuous_metrics = [
     MetricKind.Cos,
     MetricKind.L2sq,
@@ -80,6 +81,8 @@ def test_index(
     connectivity: int,
     jit: bool,
 ):
+    if jit:
+        metric = njit(ndim, metric, index_type)
     index = Index(
         metric=metric,
         ndim=ndim,
@@ -87,7 +90,6 @@ def test_index(
         connectivity=connectivity,
         expansion_add=DEFAULT_EXPANSION_ADD,
         expansion_search=DEFAULT_EXPANSION_SEARCH,
-        jit=jit,
     )
     assert index.ndim == ndim
     assert index.connectivity == connectivity
@@ -176,7 +178,10 @@ def test_index_batch(
 @pytest.mark.parametrize("index_type", index_types)
 @pytest.mark.parametrize("numpy_type", numpy_types)
 def test_exact_recall(
-    metric: MetricKind, batch_size: int, index_type: ScalarKind, numpy_type: str
+    metric: MetricKind,
+    batch_size: int,
+    index_type: ScalarKind,
+    numpy_type: str,
 ):
     ndim: int = batch_size
     index = Index(ndim=ndim, metric=metric, dtype=index_type)
@@ -436,7 +441,10 @@ def test_index_peachpy(ndim: int, batch_size: int):
 @pytest.mark.parametrize("connectivity", connectivity_options)
 @pytest.mark.parametrize("batch_size", batch_sizes)
 def test_bitwise_index(
-    bits: int, metric: MetricKind, connectivity: int, batch_size: int
+    bits: int,
+    metric: MetricKind,
+    connectivity: int,
+    batch_size: int,
 ):
     index = Index(ndim=bits, metric=metric, connectivity=connectivity)
 

package/python/usearch/index.py

@@ -167,10 +167,9 @@ class Index:
 
     def __init__(
         self,
-        ndim: int,
+        ndim: int = 0,
         metric: Union[str, MetricKind, CompiledMetric] = MetricKind.IP,
-        dtype: Optional[str] = None,
-        jit: bool = False,
+        dtype: Optional[Union[str, ScalarKind]] = None,
         connectivity: int = DEFAULT_CONNECTIVITY,
         expansion_add: int = DEFAULT_EXPANSION_ADD,
         expansion_search: int = DEFAULT_EXPANSION_SEARCH,
@@ -182,6 +181,10 @@ class Index:
 
         :param ndim: Number of vector dimensions
         :type ndim: int
+            Required for some metrics, optional for others.
+            Haversine, for example, only applies to 2-dimensional latitude/longitude
+            coordinates. Angular (Cos) and Euclidean (L2sq), obviously, apply to
+            vectors with arbitrary number of dimensions.
 
         :param metric: Distance function, defaults to MetricKind.IP
         :type metric: Union[MetricKind, Callable, str], optional
@@ -191,18 +194,12 @@ class Index:
             Not every kind is JIT-able. For Jaccard distance, use `SparseIndex`.
 
         :param dtype: Scalar type for internal vector storage, defaults to None
-        :type dtype: str, optional
+        :type dtype: Optional[Union[str, ScalarKind]], optional
             For continuous metrics can be: f16, f32, f64, or f8.
             For bitwise metrics it's implementation-defined, and can't change.
             Example: you can use the `f16` index with `f32` vectors in Euclidean space,
             which will be automatically downcasted.
 
-        :param jit: Enable Numba to JIT compile the metric, defaults to False
-        :type jit: bool, optional
-            This can result in up-to 3x performance difference on very large vectors
-            and very recent hardware, as the Python module is compiled with high
-            compatibility in mind and avoids very fancy assembly instructions.
-
         :param connectivity: Connections per node in HNSW, defaults to None
         :type connectivity: Optional[int], optional
             Hyper-parameter for the number of Graph connections
@@ -231,21 +228,6 @@ class Index:
         """
 
         metric = _normalize_metric(metric)
-        if isinstance(metric, MetricKind) and jit:
-            try:
-                from usearch.numba import jit
-            except ImportError:
-                raise ModuleNotFoundError(
-                    "To use JIT install Numba with `pip install numba`."
-                    "Alternatively, reinstall with `pip install usearch[jit]`"
-                )
-
-            metric = jit(
-                ndim=ndim,
-                metric=metric,
-                dtype=_normalize_dtype(dtype, metric),
-            )
-
         if isinstance(metric, MetricKind):
             self._metric_kind = metric
             self._metric_jit = None

package/python/usearch/numba.py

@@ -4,31 +4,19 @@
 # into the primary `Index` class, connecting USearch with Numba.
 from math import sqrt
 
-from numba import cfunc, types, carray
-
 from usearch.index import MetricKind, ScalarKind, MetricSignature, CompiledMetric
 
 
-signature_i8args = types.float32(types.CPointer(types.int8), types.CPointer(types.int8))
-
-signature_f16args = types.float32(
-    types.CPointer(types.float16), types.CPointer(types.float16)
-)
-
-signature_f32args = types.float32(
-    types.CPointer(types.float32), types.CPointer(types.float32)
-)
-
-signature_f64args = types.float32(
-    types.CPointer(types.float64), types.CPointer(types.float64)
-)
-
-
 def jit(
-    ndim: int, metric: MetricKind = MetricKind.Cos, dtype: ScalarKind = ScalarKind.F32
+    ndim: int,
+    metric: MetricKind = MetricKind.Cos,
+    dtype: ScalarKind = ScalarKind.F32,
 ) -> CompiledMetric:
-    """JIT-compiles a distance metric specifically tuned for the target hardware
-    and number of dimensions.
+    """JIT-compiles the metric for target hardware and number of dimensions.
+
+    This can result in up-to 3x performance difference on very large vectors
+    and very recent hardware, as the Python module is compiled with high
+    compatibility in mind and avoids very fancy assembly instructions.
 
     Uses Numba `cfunc` functionality, annotating it with Numba `types` instead
     of `ctypes` to support half-precision.
@@ -37,6 +25,24 @@ def jit(
     assert isinstance(metric, MetricKind)
     assert isinstance(dtype, ScalarKind)
 
+    from numba import cfunc, types, carray
+
+    signature_i8args = types.float32(
+        types.CPointer(types.int8), types.CPointer(types.int8)
+    )
+
+    signature_f16args = types.float32(
+        types.CPointer(types.float16), types.CPointer(types.float16)
+    )
+
+    signature_f32args = types.float32(
+        types.CPointer(types.float32), types.CPointer(types.float32)
+    )
+
+    signature_f64args = types.float32(
+        types.CPointer(types.float64), types.CPointer(types.float64)
+    )
+
     numba_supported_types = (
         ScalarKind.F8,
         # Half-precision is still unsupported
@@ -46,7 +52,7 @@ def jit(
         ScalarKind.F64,
     )
     if dtype not in numba_supported_types:
-        return None
+        return metric
 
     scalar_kind_to_accumulator_type = {
         ScalarKind.F8: types.int32,
@@ -109,7 +115,7 @@ def jit(
 
     pointer = cfunc(scalar_kind_to_signature[dtype])(metric_kind_to_function[metric])
     return CompiledMetric(
-        pointer=pointer,
+        pointer=pointer.address,
         kind=metric,
         signature=MetricSignature.ArrayArray,
     )