Fix distance (cpu)

Author: tomwhite

Diff for: sgkit/distance/api.py

@@ -109,9 +109,9 @@ def pairwise_distance(
             f"Invalid Device, expected one of {valid_devices}, got: {device}"
         )
     try:
-        map_func_name = f"{metric}_map_{device}"
+        map_pairwise_func_name = f"{metric}_map_pairwise_{device}"
         reduce_func_name = f"{metric}_reduce_{device}"
-        map_func = getattr(metrics, map_func_name)
+        map_pairwise_func = getattr(metrics, map_pairwise_func_name)
         reduce_func = getattr(metrics, reduce_func_name)
         n_map_param = metrics.N_MAP_PARAM[metric]
     except AttributeError:
@@ -123,24 +123,7 @@ def pairwise_distance(
     if x.ndim != 2:
         raise ValueError(f"2-dimensional array expected, got '{x.ndim}'")
 
-    # setting this variable outside of _pairwise to avoid it's recreation
-    # in every iteration, which eventually leads to increase in dask
-    # graph serialisation/deserialisation time significantly
-    metric_param = np.empty(n_map_param, dtype=x.dtype)
-
-    def _pairwise_cpu(f: ArrayLike, g: ArrayLike) -> ArrayLike:
-        result: ArrayLike = map_func(f[:, None, :], g, metric_param)
-        # Adding a new axis to help combine chunks along this axis in the
-        # reduction step (see the _aggregate and _combine functions below).
-        return result[..., np.newaxis]
-
-    def _pairwise_gpu(f: ArrayLike, g: ArrayLike) -> ArrayLike:  # pragma: no cover
-        result = map_func(f, g)
-        return result[..., np.newaxis]
-
-    pairwise_func = _pairwise_cpu
-    if device == "gpu":
-        pairwise_func = _pairwise_gpu  # pragma: no cover
+    pairwise_func = map_pairwise_func
 
     # concatenate in blockwise leads to high memory footprints, so instead
     # we perform blockwise without contraction followed by reduction.

Diff for: sgkit/distance/metrics.py

@@ -59,6 +59,14 @@ def euclidean_map_cpu(
     out[:] = square_sum
 
 
+def euclidean_map_pairwise_cpu(f: ArrayLike, g: ArrayLike) -> ArrayLike:
+    metric_param = np.empty(1, dtype=f.dtype)
+    result: ArrayLike = euclidean_map_cpu(f[:, None, :], g, metric_param)
+    # Adding a new axis to help combine chunks along this axis in the
+    # reduction step (see the _aggregate and _combine functions below).
+    return result[..., np.newaxis]
+
+
 def euclidean_reduce_cpu(v: ArrayLike) -> ArrayLike:  # pragma: no cover
     """Corresponding "reduce" function for euclidean distance.
 
@@ -138,6 +146,15 @@ def correlation_map_cpu(
     )
 
 
+def correlation_map_pairwise_cpu(f: ArrayLike, g: ArrayLike) -> ArrayLike:
+    # TODO: note that allocating this array here goes against the advice in api.py about dask graph serialization time
+    metric_param = np.empty(6, dtype=f.dtype)
+    result: ArrayLike = correlation_map_cpu(f[:, None, :], g, metric_param)
+    # Adding a new axis to help combine chunks along this axis in the
+    # reduction step (see the _aggregate and _combine functions below).
+    return result[..., np.newaxis]
+
+
 @numba_guvectorize(  # type: ignore
     [
         "void(float32[:, :], float32[:])",