Updating pylibraft pairwise_distance to cuvs (while still being backw…

…ards compatible with pylibraft when appropriate)

cupyx/scipy/spatial/distance.py

@@ -1,9 +1,16 @@
 import cupy
+
+cuvs_available = False
 try:
-    from pylibraft.distance import pairwise_distance
-    pylibraft_available = True
-except ModuleNotFoundError:
-    pylibraft_available = False
+    from cuvs.distance import pairwise_distance
+    cuvs_available = True
+except ImportError:
+    try:
+        #cuVS distance primitives were previously in pylibraft
+        from pylibraft.distance import pairwise_distance
+        cuvs_available = True
+    except ImportError:
+        cuvs_available = False
 
 
 def _convert_to_type(X, out_type):
@@ -116,8 +123,8 @@ def minkowski(u, v, p):
     Returns:
         minkowski (double): The Minkowski distance between vectors `u` and `v`.
     """
-    if not pylibraft_available:
-        raise RuntimeError('pylibraft is not installed')
+    if not cuvs_available:
+        raise RuntimeError('cuVS is not installed')
     u = cupy.asarray(u)
     v = cupy.asarray(v)
     output_arr = cupy.zeros((1,), dtype=u.dtype)
@@ -141,8 +148,8 @@ def canberra(u, v):
     Returns:
         canberra (double): The Canberra distance between vectors `u` and `v`.
     """
-    if not pylibraft_available:
-        raise RuntimeError('pylibraft is not installed')
+    if not cuvs_available:
+        raise RuntimeError('cuVS is not installed')
     u = cupy.asarray(u)
     v = cupy.asarray(v)
     output_arr = cupy.zeros((1,), dtype=u.dtype)
@@ -166,8 +173,8 @@ def chebyshev(u, v):
     Returns:
         chebyshev (double): The Chebyshev distance between vectors `u` and `v`.
     """
-    if not pylibraft_available:
-        raise RuntimeError('pylibraft is not installed')
+    if not cuvs_available:
+        raise RuntimeError('cuVS is not installed')
     u = cupy.asarray(u)
     v = cupy.asarray(v)
     output_arr = cupy.zeros((1,), dtype=u.dtype)
@@ -192,8 +199,8 @@ def cityblock(u, v):
         cityblock (double): The City Block distance between
         vectors `u` and `v`.
     """
-    if not pylibraft_available:
-        raise RuntimeError('pylibraft is not installed')
+    if not cuvs_available:
+        raise RuntimeError('cuVS is not installed')
     u = cupy.asarray(u)
     v = cupy.asarray(v)
     output_arr = cupy.zeros((1,), dtype=u.dtype)
@@ -222,8 +229,8 @@ def correlation(u, v):
         correlation (double): The correlation distance between
         vectors `u` and `v`.
     """
-    if not pylibraft_available:
-        raise RuntimeError('pylibraft is not installed')
+    if not cuvs_available:
+        raise RuntimeError('cuVS is not installed')
     u = cupy.asarray(u)
     v = cupy.asarray(v)
     output_arr = cupy.zeros((1,), dtype=u.dtype)
@@ -249,8 +256,8 @@ def cosine(u, v):
     Returns:
         cosine (double): The Cosine distance between vectors `u` and `v`.
     """
-    if not pylibraft_available:
-        raise RuntimeError('pylibraft is not installed')
+    if not cuvs_available:
+        raise RuntimeError('cuVS is not installed')
     u = cupy.asarray(u)
     v = cupy.asarray(v)
     output_arr = cupy.zeros((1,), dtype=u.dtype)
@@ -278,8 +285,8 @@ def hamming(u, v):
     Returns:
         hamming (double): The Hamming distance between vectors `u` and `v`.
     """
-    if not pylibraft_available:
-        raise RuntimeError('pylibraft is not installed')
+    if not cuvs_available:
+        raise RuntimeError('cuVS is not installed')
     u = cupy.asarray(u)
     v = cupy.asarray(v)
     output_arr = cupy.zeros((1,), dtype=u.dtype)
@@ -303,8 +310,8 @@ def euclidean(u, v):
     Returns:
         euclidean (double): The Euclidean distance between vectors `u` and `v`.
     """
-    if not pylibraft_available:
-        raise RuntimeError('pylibraft is not installed')
+    if not cuvs_available:
+        raise RuntimeError('cuVS is not installed')
     u = cupy.asarray(u)
     v = cupy.asarray(v)
     output_arr = cupy.zeros((1,), dtype=u.dtype)
@@ -332,8 +339,8 @@ def jensenshannon(u, v):
         jensenshannon (double): The Jensen-Shannon distance between
         vectors `u` and `v`.
     """
-    if not pylibraft_available:
-        raise RuntimeError('pylibraft is not installed')
+    if not cuvs_available:
+        raise RuntimeError('cuVS is not installed')
     u = cupy.asarray(u)
     v = cupy.asarray(v)
     output_arr = cupy.zeros((1,), dtype=u.dtype)
@@ -361,8 +368,8 @@ def russellrao(u, v):
     Returns:
         hamming (double): The Hamming distance between vectors `u` and `v`.
     """
-    if not pylibraft_available:
-        raise RuntimeError('pylibraft is not installed')
+    if not cuvs_available:
+        raise RuntimeError('cuVS is not installed')
     u = cupy.asarray(u)
     v = cupy.asarray(v)
     output_arr = cupy.zeros((1,), dtype=u.dtype)
@@ -387,8 +394,8 @@ def sqeuclidean(u, v):
         sqeuclidean (double): The squared Euclidean distance between
         vectors `u` and `v`.
     """
-    if not pylibraft_available:
-        raise RuntimeError('pylibraft is not installed.')
+    if not cuvs_available:
+        raise RuntimeError('cuVS is not installed.')
     u = cupy.asarray(u)
     v = cupy.asarray(v)
     output_arr = cupy.zeros((1,), dtype=u.dtype)
@@ -414,8 +421,8 @@ def hellinger(u, v):
         hellinger (double): The Hellinger distance between
         vectors `u` and `v`.
     """
-    if not pylibraft_available:
-        raise RuntimeError('pylibraft is not installed')
+    if not cuvs_available:
+        raise RuntimeError('cuVS is not installed')
     u = cupy.asarray(u)
     v = cupy.asarray(v)
     output_arr = cupy.zeros((1,), dtype=u.dtype)
@@ -440,8 +447,8 @@ def kl_divergence(u, v):
         kl_divergence (double): The Kullback-Leibler divergence between
         vectors `u` and `v`.
     """
-    if not pylibraft_available:
-        raise RuntimeError('pylibraft is not installed')
+    if not cuvs_available:
+        raise RuntimeError('cuVS is not installed')
     u = cupy.asarray(u)
     v = cupy.asarray(v)
     output_arr = cupy.zeros((1,), dtype=u.dtype)
@@ -479,10 +486,14 @@ def cdist(XA, XB, metric='euclidean', out=None, **kwargs):
             ``dist(u=XA[i], v=XB[j])`` is computed and stored in the
             :math:`ij` th entry.
     """
-    if not pylibraft_available:
-        raise RuntimeError('pylibraft is not installed')
-    XA = cupy.asarray(XA, dtype='float32')
-    XB = cupy.asarray(XB, dtype='float32')
+    if not cuvs_available:
+        raise RuntimeError('cuVS is not installed')
+
+    if XA.dtype not in ['float16', 'float32', 'float64']:
+        raise ValueError('XA must be float16, float32, or float64')
+
+    if XB.dtype not in ['float16', 'float32', 'float64']:
+        raise ValueError('XB must be float16, float32, or float64')
 
     s = XA.shape
     sB = XB.shape
@@ -501,8 +512,9 @@ def cdist(XA, XB, metric='euclidean', out=None, **kwargs):
     p = kwargs["p"] if "p" in kwargs else 2.0
 
     if out is not None:
-        if out.dtype != 'float32':
-            out = out.astype('float32', copy=False)
+        if out.dtype not in ['float16', 'float32', 'float64']:
+            raise ValueError('out must be float16, float32, or float64')
+
         if out.shape != (mA, mB):
             cupy.resize(out, (mA, mB))
         out[:] = 0.0