From 753a05821eaf5cf5d7154a1a4fb3e4221d7504d1 Mon Sep 17 00:00:00 2001
From: "Corey J. Nolet" <cjnolet@gmail.com>
Date: Thu, 10 Nov 2022 17:12:02 -0500
Subject: [PATCH 01/14] renaming to device_ndarray

---
 .../cluster/experimental/CMakeLists.txt       | 28 +++++++
 python/pylibraft/pylibraft/common/__init__.py |  1 +
 .../pylibraft/common/device_ndarray.py        | 80 +++++++++++++++++++
 .../pylibraft/test/test_device_buffer.py      | 45 +++++++++++
 .../pylibraft/pylibraft/test/test_distance.py |  6 +-
 .../pylibraft/test/test_fused_l2_argmin.py    |  8 +-
 .../pylibraft/pylibraft/test/test_kmeans.py   | 16 ++--
 .../pylibraft/pylibraft/test/test_random.py   | 10 +--
 .../pylibraft/pylibraft/testing/__init__.py   | 14 ----
 python/pylibraft/pylibraft/testing/utils.py   | 46 -----------
 10 files changed, 174 insertions(+), 80 deletions(-)
 create mode 100644 python/pylibraft/pylibraft/cluster/experimental/CMakeLists.txt
 create mode 100644 python/pylibraft/pylibraft/common/device_ndarray.py
 create mode 100644 python/pylibraft/pylibraft/test/test_device_buffer.py
 delete mode 100644 python/pylibraft/pylibraft/testing/__init__.py
 delete mode 100644 python/pylibraft/pylibraft/testing/utils.py

diff --git a/python/pylibraft/pylibraft/cluster/experimental/CMakeLists.txt b/python/pylibraft/pylibraft/cluster/experimental/CMakeLists.txt
new file mode 100644
index 0000000000..6c32e92a0f
--- /dev/null
+++ b/python/pylibraft/pylibraft/cluster/experimental/CMakeLists.txt
@@ -0,0 +1,28 @@
+# =============================================================================
+# Copyright (c) 2022, NVIDIA CORPORATION.
+#
+# Licensed under the Apache License, Version 2.0 (the "License"); you may not use this file except
+# in compliance with the License. You may obtain a copy of the License at
+#
+# http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software distributed under the License
+# is distributed on an "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express
+# or implied. See the License for the specific language governing permissions and limitations under
+# the License.
+# =============================================================================
+
+# Set the list of Cython files to build
+set(cython_sources kmeans.pyx)
+set(linked_libraries raft::raft raft::distance)
+
+# Build all of the Cython targets
+rapids_cython_create_modules(
+        CXX
+        SOURCE_FILES "${cython_sources}"
+        LINKED_LIBRARIES "${linked_libraries}"
+        MODULE_PREFIX cluster_experimental_)
+
+foreach(cython_module IN LISTS RAPIDS_CYTHON_CREATED_TARGETS)
+    set_target_properties(${cython_module} PROPERTIES INSTALL_RPATH "\$ORIGIN;\$ORIGIN/../library")
+endforeach()
diff --git a/python/pylibraft/pylibraft/common/__init__.py b/python/pylibraft/pylibraft/common/__init__.py
index 7872599a78..78c15ddcf1 100644
--- a/python/pylibraft/pylibraft/common/__init__.py
+++ b/python/pylibraft/pylibraft/common/__init__.py
@@ -15,3 +15,4 @@
 
 from .cuda import Stream
 from .handle import Handle
+from .device_ndarray import device_ndarray
diff --git a/python/pylibraft/pylibraft/common/device_ndarray.py b/python/pylibraft/pylibraft/common/device_ndarray.py
new file mode 100644
index 0000000000..5bc320ccda
--- /dev/null
+++ b/python/pylibraft/pylibraft/common/device_ndarray.py
@@ -0,0 +1,80 @@
+# Copyright (c) 2022, NVIDIA CORPORATION.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+#
+
+import numpy as np
+import rmm
+
+class device_ndarray:
+
+    def __init__(self, np_ndarray):
+        """
+        Construct a raft.device_ndarray wrapper around a numpy.ndarray (on host)
+
+        Parameters
+        ----------
+        ndarray : Any array that provides a valid __array_interface__ or __cuda_array_interface__
+        """
+        self.ndarray_ = np_ndarray
+        order = "C" if self.c_contiguous else "F"
+        self.device_buffer_ = rmm.DeviceBuffer.to_device(
+            self.ndarray_.tobytes(order=order)
+        )
+
+    @property
+    def c_contiguous(self):
+        array_interface = self.ndarray_.__array_interface__
+        strides = self.strides
+        return strides is None or \
+               array_interface["strides"][1] == self.dtype.itemsize
+
+    @property
+    def f_contiguous(self):
+        return not self.c_contiguous
+
+    @property
+    def dtype(self):
+        array_interface = self.ndarray_.__array_interface__
+        return np.dtype(array_interface["typestr"])
+
+    @property
+    def shape(self):
+        array_interface = self.ndarray_.__array_interface__
+        return array_interface["shape"]
+
+    @property
+    def strides(self):
+        array_interface = self.ndarray_.__array_interface__
+        return None if "strides" not in array_interface else \
+            array_interface["strides"]
+
+    @property
+    def __cuda_array_interface__(self):
+        device_cai = self.device_buffer_.__cuda_array_interface__
+        host_cai = self.ndarray_.__array_interface__.copy()
+        host_cai["data"] = (device_cai["data"][0], device_cai["data"][1])
+
+        return host_cai
+
+    def copy_to_host(self):
+        ret = (
+            np.frombuffer(
+                self.device_buffer_.tobytes(),
+                dtype=self.dtype,
+                like=self.ndarray_,
+            )
+            .astype(self.dtype)
+        )
+        ret = np.lib.stride_tricks.as_strided(ret, self.shape, self.strides)
+        return ret
diff --git a/python/pylibraft/pylibraft/test/test_device_buffer.py b/python/pylibraft/pylibraft/test/test_device_buffer.py
new file mode 100644
index 0000000000..6dba8c53f8
--- /dev/null
+++ b/python/pylibraft/pylibraft/test/test_device_buffer.py
@@ -0,0 +1,45 @@
+# Copyright (c) 2022, NVIDIA CORPORATION.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+#
+
+import numpy as np
+import pytest
+from pylibraft.common import device_ndarray
+
+@pytest.mark.parametrize("order", ["F", "C"])
+@pytest.mark.parametrize("dtype", [np.float32, np.float64])
+def test_basic_attributes(order, dtype):
+
+    a = np.random.random((500, 2)).astype(dtype)
+
+    if order == "C":
+        a = np.ascontiguousarray(a)
+    else:
+        a = np.asfortranarray(a)
+
+    db = device_ndarray(a)
+
+    db_host = db.copy_to_host()
+
+    assert a.shape == db.shape
+    assert a.dtype == db.dtype
+    assert a.data.f_contiguous == db.f_contiguous
+    assert a.data.f_contiguous == db_host.data.f_contiguous
+
+    print(str(a.strides))
+    print(str(db.strides))
+
+    assert a.data.c_contiguous == db.c_contiguous
+    assert a.data.c_contiguous == db_host.data.c_contiguous
+    np.testing.assert_array_equal(a.tolist(), db_host.tolist())
diff --git a/python/pylibraft/pylibraft/test/test_distance.py b/python/pylibraft/pylibraft/test/test_distance.py
index 670beb156e..b2db466aac 100644
--- a/python/pylibraft/pylibraft/test/test_distance.py
+++ b/python/pylibraft/pylibraft/test/test_distance.py
@@ -19,7 +19,7 @@
 
 from pylibraft.common import Handle
 from pylibraft.distance import pairwise_distance
-from pylibraft.testing.utils import TestDeviceBuffer
+from pylibraft.common import device_ndarray
 
 
 @pytest.mark.parametrize("n_rows", [100])
@@ -61,8 +61,8 @@ def test_distance(n_rows, n_cols, metric, order, dtype):
 
     expected[expected <= 1e-5] = 0.0
 
-    input1_device = TestDeviceBuffer(input1, order)
-    output_device = TestDeviceBuffer(output, order)
+    input1_device = device_ndarray(input1)
+    output_device = device_ndarray(output)
 
     handle = Handle()
     pairwise_distance(input1_device, input1_device, output_device, metric)
diff --git a/python/pylibraft/pylibraft/test/test_fused_l2_argmin.py b/python/pylibraft/pylibraft/test/test_fused_l2_argmin.py
index abe56f2b04..3566bcdd1a 100644
--- a/python/pylibraft/pylibraft/test/test_fused_l2_argmin.py
+++ b/python/pylibraft/pylibraft/test/test_fused_l2_argmin.py
@@ -19,7 +19,7 @@
 
 from pylibraft.common import Handle
 from pylibraft.distance import fused_l2_nn_argmin
-from pylibraft.testing.utils import TestDeviceBuffer
+from pylibraft.common import device_ndarray
 
 
 @pytest.mark.parametrize("n_rows", [10, 100])
@@ -38,9 +38,9 @@ def test_fused_l2_nn_minarg(n_rows, n_cols, n_clusters, dtype):
 
     expected = expected.argmin(axis=1)
 
-    input1_device = TestDeviceBuffer(input1, "C")
-    input2_device = TestDeviceBuffer(input2, "C")
-    output_device = TestDeviceBuffer(output, "C")
+    input1_device = device_ndarray(input1)
+    input2_device = device_ndarray(input2)
+    output_device = device_ndarray(output)
 
     handle = Handle()
     fused_l2_nn_argmin(
diff --git a/python/pylibraft/pylibraft/test/test_kmeans.py b/python/pylibraft/pylibraft/test/test_kmeans.py
index d198ac2f8f..547846af4a 100644
--- a/python/pylibraft/pylibraft/test/test_kmeans.py
+++ b/python/pylibraft/pylibraft/test/test_kmeans.py
@@ -19,7 +19,7 @@
 from pylibraft.cluster.kmeans import compute_new_centroids
 from pylibraft.common import Handle
 from pylibraft.distance import pairwise_distance
-from pylibraft.testing.utils import TestDeviceBuffer
+from pylibraft.common import device_ndarray
 
 
 @pytest.mark.parametrize("n_rows", [100])
@@ -39,34 +39,34 @@ def test_compute_new_centroids(
     handle = Handle()
 
     X = np.random.random_sample((n_rows, n_cols)).astype(dtype)
-    X_device = TestDeviceBuffer(X, order)
+    X_device = device_ndarray(X)
 
     centroids = X[:n_clusters]
-    centroids_device = TestDeviceBuffer(centroids, order)
+    centroids_device = device_ndarray(centroids)
 
     weight_per_cluster = np.zeros((n_clusters,), dtype=dtype)
     weight_per_cluster_device = (
-        TestDeviceBuffer(weight_per_cluster, order)
+        device_ndarray(weight_per_cluster)
         if additional_args
         else None
     )
 
     new_centroids = np.zeros((n_clusters, n_cols), dtype=dtype)
-    new_centroids_device = TestDeviceBuffer(new_centroids, order)
+    new_centroids_device = device_ndarray(new_centroids)
 
     sample_weights = np.ones((n_rows,)).astype(dtype) / n_rows
     sample_weights_device = (
-        TestDeviceBuffer(sample_weights, order) if additional_args else None
+        device_ndarray(sample_weights) if additional_args else None
     )
 
     # Compute new centroids naively
     dists = np.zeros((n_rows, n_clusters), dtype=dtype)
-    dists_device = TestDeviceBuffer(dists, order)
+    dists_device = device_ndarray(dists)
     pairwise_distance(X_device, centroids_device, dists_device, metric=metric)
     handle.sync()
 
     labels = np.argmin(dists_device.copy_to_host(), axis=1).astype(np.int32)
-    labels_device = TestDeviceBuffer(labels, order)
+    labels_device = device_ndarray(labels)
 
     expected_centers = np.empty((n_clusters, n_cols), dtype=dtype)
     expected_wX = X * sample_weights.reshape((-1, 1))
diff --git a/python/pylibraft/pylibraft/test/test_random.py b/python/pylibraft/pylibraft/test/test_random.py
index 77494ea277..8048513234 100644
--- a/python/pylibraft/pylibraft/test/test_random.py
+++ b/python/pylibraft/pylibraft/test/test_random.py
@@ -18,7 +18,7 @@
 
 from pylibraft.common import Handle
 from pylibraft.random import rmat
-from pylibraft.testing.utils import TestDeviceBuffer
+from pylibraft.common import device_ndarray
 
 
 def generate_theta(r_scale, c_scale):
@@ -34,7 +34,7 @@ def generate_theta(r_scale, c_scale):
         theta[4 * i + 1] = b / total
         theta[4 * i + 2] = c / total
         theta[4 * i + 3] = d / total
-    theta_device = TestDeviceBuffer(theta, "C")
+    theta_device = device_ndarray(theta)
     return theta, theta_device
 
 
@@ -45,7 +45,7 @@ def generate_theta(r_scale, c_scale):
 def test_rmat(n_edges, r_scale, c_scale, dtype):
     theta, theta_device = generate_theta(r_scale, c_scale)
     out_buff = np.empty((n_edges, 2), dtype=dtype)
-    output_device = TestDeviceBuffer(out_buff, "C")
+    output_device = device_ndarray(out_buff)
 
     handle = Handle()
     rmat(output_device, theta_device, r_scale, c_scale, 12345, handle=handle)
@@ -68,7 +68,7 @@ def test_rmat_exception():
     dtype = np.int32
     with pytest.raises(Exception) as exception:
         out_buff = np.empty((n_edges, 2), dtype=dtype)
-        output_device = TestDeviceBuffer(out_buff, "C")
+        output_device = device_ndarray(out_buff)
         rmat(output_device, None, r_scale, c_scale, 12345)
         assert exception is not None
         assert exception.message == "'theta' cannot be None!"
@@ -84,7 +84,7 @@ def test_rmat_valueerror():
     r_scale = c_scale = 16
     with pytest.raises(ValueError) as exception:
         out_buff = np.empty((n_edges, 2), dtype=np.int16)
-        output_device = TestDeviceBuffer(out_buff, "C")
+        output_device = device_ndarray(out_buff)
         theta, theta_device = generate_theta(r_scale, c_scale)
         rmat(output_device, theta_device, r_scale, c_scale, 12345)
         assert exception is not None
diff --git a/python/pylibraft/pylibraft/testing/__init__.py b/python/pylibraft/pylibraft/testing/__init__.py
deleted file mode 100644
index 273b4497cc..0000000000
--- a/python/pylibraft/pylibraft/testing/__init__.py
+++ /dev/null
@@ -1,14 +0,0 @@
-# Copyright (c) 2022, NVIDIA CORPORATION.
-#
-# Licensed under the Apache License, Version 2.0 (the "License");
-# you may not use this file except in compliance with the License.
-# You may obtain a copy of the License at
-#
-#     http://www.apache.org/licenses/LICENSE-2.0
-#
-# Unless required by applicable law or agreed to in writing, software
-# distributed under the License is distributed on an "AS IS" BASIS,
-# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
-# See the License for the specific language governing permissions and
-# limitations under the License.
-#
diff --git a/python/pylibraft/pylibraft/testing/utils.py b/python/pylibraft/pylibraft/testing/utils.py
deleted file mode 100644
index 86cf4558db..0000000000
--- a/python/pylibraft/pylibraft/testing/utils.py
+++ /dev/null
@@ -1,46 +0,0 @@
-# Copyright (c) 2022, NVIDIA CORPORATION.
-#
-# Licensed under the Apache License, Version 2.0 (the "License");
-# you may not use this file except in compliance with the License.
-# You may obtain a copy of the License at
-#
-#     http://www.apache.org/licenses/LICENSE-2.0
-#
-# Unless required by applicable law or agreed to in writing, software
-# distributed under the License is distributed on an "AS IS" BASIS,
-# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
-# See the License for the specific language governing permissions and
-# limitations under the License.
-#
-
-import numpy as np
-
-import rmm
-
-
-class TestDeviceBuffer:
-    def __init__(self, ndarray, order):
-
-        self.ndarray_ = ndarray
-        self.device_buffer_ = rmm.DeviceBuffer.to_device(
-            ndarray.ravel(order=order).tobytes()
-        )
-
-    @property
-    def __cuda_array_interface__(self):
-        device_cai = self.device_buffer_.__cuda_array_interface__
-        host_cai = self.ndarray_.__array_interface__.copy()
-        host_cai["data"] = (device_cai["data"][0], device_cai["data"][1])
-
-        return host_cai
-
-    def copy_to_host(self):
-        return (
-            np.frombuffer(
-                self.device_buffer_.tobytes(),
-                dtype=self.ndarray_.dtype,
-                like=self.ndarray_,
-            )
-            .astype(self.ndarray_.dtype)
-            .reshape(self.ndarray_.shape)
-        )

From 2bd251ae2bf1abc8c0245036e214949a1fcfb98b Mon Sep 17 00:00:00 2001
From: "Corey J. Nolet" <cjnolet@gmail.com>
Date: Thu, 10 Nov 2022 17:13:46 -0500
Subject: [PATCH 02/14] Using pre-commit

---
 python/pylibraft/pylibraft/common/__init__.py | 3 ++-
 1 file changed, 2 insertions(+), 1 deletion(-)

diff --git a/python/pylibraft/pylibraft/common/__init__.py b/python/pylibraft/pylibraft/common/__init__.py
index 78c15ddcf1..4c6f0d686a 100644
--- a/python/pylibraft/pylibraft/common/__init__.py
+++ b/python/pylibraft/pylibraft/common/__init__.py
@@ -13,6 +13,7 @@
 # limitations under the License.
 #
 
+
 from .cuda import Stream
-from .handle import Handle
 from .device_ndarray import device_ndarray
+from .handle import Handle

From ac561ecb50c3b8970aac03756a2a818e9c449f01 Mon Sep 17 00:00:00 2001
From: "Corey J. Nolet" <cjnolet@gmail.com>
Date: Thu, 10 Nov 2022 17:22:33 -0500
Subject: [PATCH 03/14] Making pairwise distance output optional

---
 .../pylibraft/distance/pairwise_distance.pyx  | 28 +++++++++++++++----
 .../pylibraft/test/test_device_buffer.py      |  7 ++---
 .../pylibraft/pylibraft/test/test_distance.py | 14 ++++++----
 3 files changed, 34 insertions(+), 15 deletions(-)

diff --git a/python/pylibraft/pylibraft/distance/pairwise_distance.pyx b/python/pylibraft/pylibraft/distance/pairwise_distance.pyx
index 76cdf0b2d3..549364ec4e 100644
--- a/python/pylibraft/pylibraft/distance/pairwise_distance.pyx
+++ b/python/pylibraft/pylibraft/distance/pairwise_distance.pyx
@@ -28,8 +28,11 @@ from .distance_type cimport DistanceType
 
 from pylibraft.common import Handle
 from pylibraft.common.handle import auto_sync_handle
+
 from pylibraft.common.handle cimport handle_t
 
+from pylibraft.common import device_ndarray
+
 
 def is_c_cont(cai, dt):
     return "strides" not in cai or \
@@ -92,7 +95,7 @@ SUPPORTED_DISTANCES = ["euclidean", "l1", "cityblock", "l2", "inner_product",
 
 
 @auto_sync_handle
-def distance(X, Y, dists, metric="euclidean", p=2.0, handle=None):
+def distance(X, Y, out=None, metric="euclidean", p=2.0, handle=None):
     """
     Compute pairwise distances between X and Y
 
@@ -107,11 +110,16 @@ def distance(X, Y, dists, metric="euclidean", p=2.0, handle=None):
 
     X : CUDA array interface compliant matrix shape (m, k)
     Y : CUDA array interface compliant matrix shape (n, k)
-    dists : Writable CUDA array interface matrix shape (m, n)
+    out : Optional writable CUDA array interface matrix shape (m, n)
     metric : string denoting the metric type (default="euclidean")
     p : metric parameter (currently used only for "minkowski")
     {handle_docstring}
 
+    Returns
+    -------
+
+    raft.device_ndarray containing pairwise distances
+
     Examples
     --------
 
@@ -144,14 +152,24 @@ def distance(X, Y, dists, metric="euclidean", p=2.0, handle=None):
 
     x_cai = X.__cuda_array_interface__
     y_cai = Y.__cuda_array_interface__
-    dists_cai = dists.__cuda_array_interface__
 
     m = x_cai["shape"][0]
     n = y_cai["shape"][0]
 
+    x_dt = np.dtype(x_cai["typestr"])
+    y_dt = np.dtype(y_cai["typestr"])
+
+    if out is None:
+        out_host = np.empty((m, n), dtype=y_dt)
+        dists = device_ndarray(out_host)
+    else:
+        dists = out
+
     x_k = x_cai["shape"][1]
     y_k = y_cai["shape"][1]
 
+    dists_cai = dists.__cuda_array_interface__
+
     if x_k != y_k:
         raise ValueError("Inputs must have same number of columns. "
                          "a=%s, b=%s" % (x_k, y_k))
@@ -163,8 +181,6 @@ def distance(X, Y, dists, metric="euclidean", p=2.0, handle=None):
     handle = handle if handle is not None else Handle()
     cdef handle_t *h = <handle_t*><size_t>handle.getHandle()
 
-    x_dt = np.dtype(x_cai["typestr"])
-    y_dt = np.dtype(y_cai["typestr"])
     d_dt = np.dtype(dists_cai["typestr"])
 
     x_c_contiguous = is_c_cont(x_cai, x_dt)
@@ -205,3 +221,5 @@ def distance(X, Y, dists, metric="euclidean", p=2.0, handle=None):
                           <float>p)
     else:
         raise ValueError("dtype %s not supported" % x_dt)
+
+    return dists
diff --git a/python/pylibraft/pylibraft/test/test_device_buffer.py b/python/pylibraft/pylibraft/test/test_device_buffer.py
index 6dba8c53f8..360eef64d2 100644
--- a/python/pylibraft/pylibraft/test/test_device_buffer.py
+++ b/python/pylibraft/pylibraft/test/test_device_buffer.py
@@ -15,8 +15,10 @@
 
 import numpy as np
 import pytest
+
 from pylibraft.common import device_ndarray
 
+
 @pytest.mark.parametrize("order", ["F", "C"])
 @pytest.mark.parametrize("dtype", [np.float32, np.float64])
 def test_basic_attributes(order, dtype):
@@ -29,17 +31,12 @@ def test_basic_attributes(order, dtype):
         a = np.asfortranarray(a)
 
     db = device_ndarray(a)
-
     db_host = db.copy_to_host()
 
     assert a.shape == db.shape
     assert a.dtype == db.dtype
     assert a.data.f_contiguous == db.f_contiguous
     assert a.data.f_contiguous == db_host.data.f_contiguous
-
-    print(str(a.strides))
-    print(str(db.strides))
-
     assert a.data.c_contiguous == db.c_contiguous
     assert a.data.c_contiguous == db_host.data.c_contiguous
     np.testing.assert_array_equal(a.tolist(), db_host.tolist())
diff --git a/python/pylibraft/pylibraft/test/test_distance.py b/python/pylibraft/pylibraft/test/test_distance.py
index b2db466aac..a08656d3aa 100644
--- a/python/pylibraft/pylibraft/test/test_distance.py
+++ b/python/pylibraft/pylibraft/test/test_distance.py
@@ -17,9 +17,8 @@
 import pytest
 from scipy.spatial.distance import cdist
 
-from pylibraft.common import Handle
+from pylibraft.common import Handle, device_ndarray
 from pylibraft.distance import pairwise_distance
-from pylibraft.common import device_ndarray
 
 
 @pytest.mark.parametrize("n_rows", [100])
@@ -39,9 +38,10 @@
         "sqeuclidean",
     ],
 )
+@pytest.mark.parametrize("inplace", [True, False])
 @pytest.mark.parametrize("order", ["F", "C"])
 @pytest.mark.parametrize("dtype", [np.float32, np.float64])
-def test_distance(n_rows, n_cols, metric, order, dtype):
+def test_distance(n_rows, n_cols, inplace, metric, order, dtype):
     input1 = np.random.random_sample((n_rows, n_cols))
     input1 = np.asarray(input1, order=order).astype(dtype)
 
@@ -62,12 +62,16 @@ def test_distance(n_rows, n_cols, metric, order, dtype):
     expected[expected <= 1e-5] = 0.0
 
     input1_device = device_ndarray(input1)
-    output_device = device_ndarray(output)
+    output_device = device_ndarray(output) if inplace else None
 
     handle = Handle()
-    pairwise_distance(input1_device, input1_device, output_device, metric)
+    ret_output = pairwise_distance(
+        input1_device, input1_device, output_device, metric
+    )
     handle.sync()
 
+    output_device = ret_output if not inplace else output_device
+
     actual = output_device.copy_to_host()
 
     actual[actual <= 1e-5] = 0.0

From f5c4a00e3ff832c311be5dd7309dd2a5af82a28d Mon Sep 17 00:00:00 2001
From: "Corey J. Nolet" <cjnolet@gmail.com>
Date: Thu, 10 Nov 2022 17:33:01 -0500
Subject: [PATCH 04/14] Updating docs for device_ndarray

---
 .../pylibraft/common/device_ndarray.py        | 60 ++++++++++++++-----
 1 file changed, 46 insertions(+), 14 deletions(-)

diff --git a/python/pylibraft/pylibraft/common/device_ndarray.py b/python/pylibraft/pylibraft/common/device_ndarray.py
index 5bc320ccda..4c43e1b555 100644
--- a/python/pylibraft/pylibraft/common/device_ndarray.py
+++ b/python/pylibraft/pylibraft/common/device_ndarray.py
@@ -14,17 +14,23 @@
 #
 
 import numpy as np
+
 import rmm
 
+
 class device_ndarray:
+    """
+    raft.device_ndarray is meant to be a very lightweight
+    __cuda_array_interface__ wrapper around a numpy.ndarray.
+    """
 
     def __init__(self, np_ndarray):
         """
-        Construct a raft.device_ndarray wrapper around a numpy.ndarray (on host)
+        Construct a raft.device_ndarray wrapper around a numpy.ndarray
 
         Parameters
         ----------
-        ndarray : Any array that provides a valid __array_interface__ or __cuda_array_interface__
+        ndarray : A numpy.ndarray which will be copied and moved to the device
         """
         self.ndarray_ = np_ndarray
         order = "C" if self.c_contiguous else "F"
@@ -34,33 +40,58 @@ def __init__(self, np_ndarray):
 
     @property
     def c_contiguous(self):
+        """
+        Is the current device_ndarray laid out in row-major format?
+        """
         array_interface = self.ndarray_.__array_interface__
         strides = self.strides
-        return strides is None or \
-               array_interface["strides"][1] == self.dtype.itemsize
+        return (
+            strides is None
+            or array_interface["strides"][1] == self.dtype.itemsize
+        )
 
     @property
     def f_contiguous(self):
+        """
+        Is the current device_ndarray laid out in column-major format?
+        """
         return not self.c_contiguous
 
     @property
     def dtype(self):
+        """
+        Datatype of the current device_ndarray instance
+        """
         array_interface = self.ndarray_.__array_interface__
         return np.dtype(array_interface["typestr"])
 
     @property
     def shape(self):
+        """
+        Shape of the current device_ndarray instance
+        """
         array_interface = self.ndarray_.__array_interface__
         return array_interface["shape"]
 
     @property
     def strides(self):
+        """
+        Strides of the current device_ndarray instance
+        """
         array_interface = self.ndarray_.__array_interface__
-        return None if "strides" not in array_interface else \
-            array_interface["strides"]
+        return (
+            None
+            if "strides" not in array_interface
+            else array_interface["strides"]
+        )
 
     @property
     def __cuda_array_interface__(self):
+        """
+        Returns the __cuda_array_interface__ compliant dict for
+        integrating with other device-enabled libraries using
+        zero-copy semantics.
+        """
         device_cai = self.device_buffer_.__cuda_array_interface__
         host_cai = self.ndarray_.__array_interface__.copy()
         host_cai["data"] = (device_cai["data"][0], device_cai["data"][1])
@@ -68,13 +99,14 @@ def __cuda_array_interface__(self):
         return host_cai
 
     def copy_to_host(self):
-        ret = (
-            np.frombuffer(
-                self.device_buffer_.tobytes(),
-                dtype=self.dtype,
-                like=self.ndarray_,
-            )
-            .astype(self.dtype)
-        )
+        """
+        Returns a new numpy.ndarray object on host with the current contents of
+        this device_ndarray
+        """
+        ret = np.frombuffer(
+            self.device_buffer_.tobytes(),
+            dtype=self.dtype,
+            like=self.ndarray_,
+        ).astype(self.dtype)
         ret = np.lib.stride_tricks.as_strided(ret, self.shape, self.strides)
         return ret

From 3d0114b0b889538986017563b023556474006922 Mon Sep 17 00:00:00 2001
From: "Corey J. Nolet" <cjnolet@gmail.com>
Date: Thu, 10 Nov 2022 17:43:22 -0500
Subject: [PATCH 05/14] Updating style

---
 .../pylibraft/cluster/experimental/CMakeLists.txt      | 10 +++++-----
 .../pylibraft/pylibraft/test/test_fused_l2_argmin.py   |  3 +--
 python/pylibraft/pylibraft/test/test_kmeans.py         |  9 ++-------
 python/pylibraft/pylibraft/test/test_random.py         |  3 +--
 4 files changed, 9 insertions(+), 16 deletions(-)

diff --git a/python/pylibraft/pylibraft/cluster/experimental/CMakeLists.txt b/python/pylibraft/pylibraft/cluster/experimental/CMakeLists.txt
index 6c32e92a0f..f2b6dca851 100644
--- a/python/pylibraft/pylibraft/cluster/experimental/CMakeLists.txt
+++ b/python/pylibraft/pylibraft/cluster/experimental/CMakeLists.txt
@@ -18,11 +18,11 @@ set(linked_libraries raft::raft raft::distance)
 
 # Build all of the Cython targets
 rapids_cython_create_modules(
-        CXX
-        SOURCE_FILES "${cython_sources}"
-        LINKED_LIBRARIES "${linked_libraries}"
-        MODULE_PREFIX cluster_experimental_)
+  CXX
+  SOURCE_FILES "${cython_sources}"
+  LINKED_LIBRARIES "${linked_libraries}" MODULE_PREFIX cluster_experimental_
+)
 
 foreach(cython_module IN LISTS RAPIDS_CYTHON_CREATED_TARGETS)
-    set_target_properties(${cython_module} PROPERTIES INSTALL_RPATH "\$ORIGIN;\$ORIGIN/../library")
+  set_target_properties(${cython_module} PROPERTIES INSTALL_RPATH "\$ORIGIN;\$ORIGIN/../library")
 endforeach()
diff --git a/python/pylibraft/pylibraft/test/test_fused_l2_argmin.py b/python/pylibraft/pylibraft/test/test_fused_l2_argmin.py
index 3566bcdd1a..cc54389a31 100644
--- a/python/pylibraft/pylibraft/test/test_fused_l2_argmin.py
+++ b/python/pylibraft/pylibraft/test/test_fused_l2_argmin.py
@@ -17,9 +17,8 @@
 import pytest
 from scipy.spatial.distance import cdist
 
-from pylibraft.common import Handle
+from pylibraft.common import Handle, device_ndarray
 from pylibraft.distance import fused_l2_nn_argmin
-from pylibraft.common import device_ndarray
 
 
 @pytest.mark.parametrize("n_rows", [10, 100])
diff --git a/python/pylibraft/pylibraft/test/test_kmeans.py b/python/pylibraft/pylibraft/test/test_kmeans.py
index 547846af4a..58028e90e8 100644
--- a/python/pylibraft/pylibraft/test/test_kmeans.py
+++ b/python/pylibraft/pylibraft/test/test_kmeans.py
@@ -17,9 +17,8 @@
 import pytest
 
 from pylibraft.cluster.kmeans import compute_new_centroids
-from pylibraft.common import Handle
+from pylibraft.common import Handle, device_ndarray
 from pylibraft.distance import pairwise_distance
-from pylibraft.common import device_ndarray
 
 
 @pytest.mark.parametrize("n_rows", [100])
@@ -32,8 +31,6 @@ def test_compute_new_centroids(
     n_rows, n_cols, metric, n_clusters, dtype, additional_args
 ):
 
-    order = "C"
-
     # A single RAFT handle can optionally be reused across
     # pylibraft functions.
     handle = Handle()
@@ -46,9 +43,7 @@ def test_compute_new_centroids(
 
     weight_per_cluster = np.zeros((n_clusters,), dtype=dtype)
     weight_per_cluster_device = (
-        device_ndarray(weight_per_cluster)
-        if additional_args
-        else None
+        device_ndarray(weight_per_cluster) if additional_args else None
     )
 
     new_centroids = np.zeros((n_clusters, n_cols), dtype=dtype)
diff --git a/python/pylibraft/pylibraft/test/test_random.py b/python/pylibraft/pylibraft/test/test_random.py
index 8048513234..229baffff5 100644
--- a/python/pylibraft/pylibraft/test/test_random.py
+++ b/python/pylibraft/pylibraft/test/test_random.py
@@ -16,9 +16,8 @@
 import numpy as np
 import pytest
 
-from pylibraft.common import Handle
+from pylibraft.common import Handle, device_ndarray
 from pylibraft.random import rmat
-from pylibraft.common import device_ndarray
 
 
 def generate_theta(r_scale, c_scale):

From a6c9d359ef9acb78209fa5d4ceec58627a2fda05 Mon Sep 17 00:00:00 2001
From: "Corey J. Nolet" <cjnolet@gmail.com>
Date: Mon, 14 Nov 2022 13:06:28 -0500
Subject: [PATCH 06/14] Adding device_ndarray.empty() staticmethod

---
 .../pylibraft/common/device_ndarray.py        | 20 ++++++++++++++++++
 ...evice_buffer.py => test_device_ndarray.py} | 21 +++++++++++++++++++
 2 files changed, 41 insertions(+)
 rename python/pylibraft/pylibraft/test/{test_device_buffer.py => test_device_ndarray.py} (67%)

diff --git a/python/pylibraft/pylibraft/common/device_ndarray.py b/python/pylibraft/pylibraft/common/device_ndarray.py
index 4c43e1b555..162172e86b 100644
--- a/python/pylibraft/pylibraft/common/device_ndarray.py
+++ b/python/pylibraft/pylibraft/common/device_ndarray.py
@@ -38,6 +38,26 @@ def __init__(self, np_ndarray):
             self.ndarray_.tobytes(order=order)
         )
 
+    @staticmethod
+    def empty(shape, dtype=np.float32, order="C"):
+        """
+        Return a new device_ndarray of given shape and type, without
+        initializing entries.
+
+        Parameters
+        ----------
+        shape : int or tuple of int
+                Shape of the empty array, e.g., (2, 3) or 2.
+        dtype : data-type, optional
+                Desired output data-type for the array, e.g, numpy.int8.
+                Default is numpy.float32.
+        order : {'C', 'F'}, optional (default: 'C')
+                Whether to store multi-dimensional dat ain row-major (C-style)
+                or column-major (Fortran-style) order in memory
+        """
+        arr = np.empty(shape, dtype=dtype, order=order)
+        return device_ndarray(arr)
+
     @property
     def c_contiguous(self):
         """
diff --git a/python/pylibraft/pylibraft/test/test_device_buffer.py b/python/pylibraft/pylibraft/test/test_device_ndarray.py
similarity index 67%
rename from python/pylibraft/pylibraft/test/test_device_buffer.py
rename to python/pylibraft/pylibraft/test/test_device_ndarray.py
index 360eef64d2..ee96abe049 100644
--- a/python/pylibraft/pylibraft/test/test_device_buffer.py
+++ b/python/pylibraft/pylibraft/test/test_device_ndarray.py
@@ -40,3 +40,24 @@ def test_basic_attributes(order, dtype):
     assert a.data.c_contiguous == db.c_contiguous
     assert a.data.c_contiguous == db_host.data.c_contiguous
     np.testing.assert_array_equal(a.tolist(), db_host.tolist())
+
+
+@pytest.mark.parametrize("order", ["F", "C"])
+@pytest.mark.parametrize("dtype", [np.float32, np.float64])
+def test_empty(order, dtype):
+
+    a = np.random.random((500, 2)).astype(dtype)
+    if order == "C":
+        a = np.ascontiguousarray(a)
+    else:
+        a = np.asfortranarray(a)
+
+    db = device_ndarray.empty(a.shape, dtype=dtype, order=order)
+    db_host = db.copy_to_host()
+
+    assert a.shape == db.shape
+    assert a.dtype == db.dtype
+    assert a.data.f_contiguous == db.f_contiguous
+    assert a.data.f_contiguous == db_host.data.f_contiguous
+    assert a.data.c_contiguous == db.c_contiguous
+    assert a.data.c_contiguous == db_host.data.c_contiguous

From 340300f65b2df92bc22796a08c7a8c1c2ba385c7 Mon Sep 17 00:00:00 2001
From: "Corey J. Nolet" <cjnolet@gmail.com>
Date: Mon, 14 Nov 2022 13:12:32 -0500
Subject: [PATCH 07/14] Using pylibraft.device_ndarray in pairwise_distances

---
 .../pylibraft/common/device_ndarray.py        | 19 +++++++++++++++++--
 .../pylibraft/distance/pairwise_distance.pyx  |  3 +--
 2 files changed, 18 insertions(+), 4 deletions(-)

diff --git a/python/pylibraft/pylibraft/common/device_ndarray.py b/python/pylibraft/pylibraft/common/device_ndarray.py
index 162172e86b..d27c567934 100644
--- a/python/pylibraft/pylibraft/common/device_ndarray.py
+++ b/python/pylibraft/pylibraft/common/device_ndarray.py
@@ -20,17 +20,32 @@
 
 class device_ndarray:
     """
-    raft.device_ndarray is meant to be a very lightweight
+    pylibraft.device_ndarray is meant to be a very lightweight
     __cuda_array_interface__ wrapper around a numpy.ndarray.
     """
 
     def __init__(self, np_ndarray):
         """
-        Construct a raft.device_ndarray wrapper around a numpy.ndarray
+        Construct a pylibraft.device_ndarray wrapper around a numpy.ndarray
 
         Parameters
         ----------
         ndarray : A numpy.ndarray which will be copied and moved to the device
+
+        Examples
+        --------
+        The device_ndarray is __cuda_array_interface__ compliant so it is
+        interoperable with other libraries that also support it, such as
+        CuPy and PyTorch. The following usage example demonstrates
+        converting a pylibraft.device_ndarray to a cupy.ndarray:
+
+        .. code-block:: python
+
+            import cupy as cp
+            from pylibraft import device_ndarray
+
+            raft_array = device_ndarray.empty((100, 50))
+            cupy_array = cp.asarray(raft_array)
         """
         self.ndarray_ = np_ndarray
         order = "C" if self.c_contiguous else "F"
diff --git a/python/pylibraft/pylibraft/distance/pairwise_distance.pyx b/python/pylibraft/pylibraft/distance/pairwise_distance.pyx
index 549364ec4e..9c1780ddc0 100644
--- a/python/pylibraft/pylibraft/distance/pairwise_distance.pyx
+++ b/python/pylibraft/pylibraft/distance/pairwise_distance.pyx
@@ -160,8 +160,7 @@ def distance(X, Y, out=None, metric="euclidean", p=2.0, handle=None):
     y_dt = np.dtype(y_cai["typestr"])
 
     if out is None:
-        out_host = np.empty((m, n), dtype=y_dt)
-        dists = device_ndarray(out_host)
+        dists = device_ndarray.empty((m, n), dtype=y_dt)
     else:
         dists = out
 

From 9d8d2b9a209e1b41c6482be15dd98e279c4390da Mon Sep 17 00:00:00 2001
From: "Corey J. Nolet" <cjnolet@gmail.com>
Date: Mon, 14 Nov 2022 13:14:02 -0500
Subject: [PATCH 08/14] Removing acciddentally checked-in file

---
 .../cluster/experimental/CMakeLists.txt       | 28 -------------------
 1 file changed, 28 deletions(-)
 delete mode 100644 python/pylibraft/pylibraft/cluster/experimental/CMakeLists.txt

diff --git a/python/pylibraft/pylibraft/cluster/experimental/CMakeLists.txt b/python/pylibraft/pylibraft/cluster/experimental/CMakeLists.txt
deleted file mode 100644
index f2b6dca851..0000000000
--- a/python/pylibraft/pylibraft/cluster/experimental/CMakeLists.txt
+++ /dev/null
@@ -1,28 +0,0 @@
-# =============================================================================
-# Copyright (c) 2022, NVIDIA CORPORATION.
-#
-# Licensed under the Apache License, Version 2.0 (the "License"); you may not use this file except
-# in compliance with the License. You may obtain a copy of the License at
-#
-# http://www.apache.org/licenses/LICENSE-2.0
-#
-# Unless required by applicable law or agreed to in writing, software distributed under the License
-# is distributed on an "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express
-# or implied. See the License for the specific language governing permissions and limitations under
-# the License.
-# =============================================================================
-
-# Set the list of Cython files to build
-set(cython_sources kmeans.pyx)
-set(linked_libraries raft::raft raft::distance)
-
-# Build all of the Cython targets
-rapids_cython_create_modules(
-  CXX
-  SOURCE_FILES "${cython_sources}"
-  LINKED_LIBRARIES "${linked_libraries}" MODULE_PREFIX cluster_experimental_
-)
-
-foreach(cython_module IN LISTS RAPIDS_CYTHON_CREATED_TARGETS)
-  set_target_properties(${cython_module} PROPERTIES INSTALL_RPATH "\$ORIGIN;\$ORIGIN/../library")
-endforeach()

From 45e521206f152517ecfbc1c1d73fccdf68299f67 Mon Sep 17 00:00:00 2001
From: "Corey J. Nolet" <cjnolet@users.noreply.github.com>
Date: Mon, 14 Nov 2022 13:29:10 -0500
Subject: [PATCH 09/14] Update
 python/pylibraft/pylibraft/common/device_ndarray.py

Co-authored-by: Ben Frederickson <github@benfrederickson.com>
---
 python/pylibraft/pylibraft/common/device_ndarray.py | 6 +++---
 1 file changed, 3 insertions(+), 3 deletions(-)

diff --git a/python/pylibraft/pylibraft/common/device_ndarray.py b/python/pylibraft/pylibraft/common/device_ndarray.py
index d27c567934..6eee4ef04d 100644
--- a/python/pylibraft/pylibraft/common/device_ndarray.py
+++ b/python/pylibraft/pylibraft/common/device_ndarray.py
@@ -53,8 +53,8 @@ def __init__(self, np_ndarray):
             self.ndarray_.tobytes(order=order)
         )
 
-    @staticmethod
-    def empty(shape, dtype=np.float32, order="C"):
+    @classmethod
+    def empty(cls, shape, dtype=np.float32, order="C"):
         """
         Return a new device_ndarray of given shape and type, without
         initializing entries.
@@ -71,7 +71,7 @@ def empty(shape, dtype=np.float32, order="C"):
                 or column-major (Fortran-style) order in memory
         """
         arr = np.empty(shape, dtype=dtype, order=order)
-        return device_ndarray(arr)
+        return cls(arr)
 
     @property
     def c_contiguous(self):

From 3f8dcc43f69c84294265f235c2c03fa1990c5846 Mon Sep 17 00:00:00 2001
From: "Corey J. Nolet" <cjnolet@gmail.com>
Date: Mon, 14 Nov 2022 14:00:20 -0500
Subject: [PATCH 10/14] Updating usage examples in readme

---
 README.md | 67 +++++++++++++++++++++++++++++++++++++++++++++++++++++--
 1 file changed, 65 insertions(+), 2 deletions(-)

diff --git a/README.md b/README.md
index ddaf8b3f8d..e0938638d1 100755
--- a/README.md
+++ b/README.md
@@ -77,11 +77,73 @@ auto metric = raft::distance::DistanceType::L2SqrtExpanded;
 raft::distance::pairwise_distance(handle, input.view(), input.view(), output.view(), metric);
 ```
 
+It's also possible to create `raft::device_mdspan` views to invoke the same API with raw pointers and shape information:
+
+```c++
+#include <raft/core/handle.hpp>
+#include <raft/core/device_mdspan.hpp>
+#include <raft/random/make_blobs.cuh>
+#include <raft/distance/distance.cuh>
+
+raft::handle_t handle;
+
+int n_samples = 5000;
+int n_features = 50;
+
+float *input;
+int *labels;
+float *output;
+
+...
+// Allocate input, labels, and output pointers
+...
+
+auto input_view = raft::make_device_matrix_view(input, n_samples, n_features);
+auto labels_view = raft::make_device_vector_view(labels, n_samples);
+auto output_view = raft::make_device_matrix_view(output, n_samples, n_samples);
+
+raft::random::make_blobs(handle, input_view, labels_view);
+
+auto metric = raft::distance::DistanceType::L2SqrtExpanded;
+raft::distance::pairwise_distance(handle, input_view, input_view, output_view, metric);
+```
+
+
 ### Python Example
 
 The `pylibraft` package contains a Python API for RAFT algorithms and primitives. `pylibraft` integrates nicely into other libraries by being very lightweight with minimal dependencies and accepting any object that supports the `__cuda_array_interface__`, such as [CuPy's ndarray](https://docs.cupy.dev/en/stable/user_guide/interoperability.html#rmm). The number of RAFT algorithms exposed in this package is continuing to grow from release to release.
 
-The example below demonstrates computing the pairwise Euclidean distances between CuPy arrays. `pylibraft` is a low-level API that prioritizes efficiency and simplicity over being pythonic, which is shown here by pre-allocating the output memory before invoking the `pairwise_distance` function. Note that CuPy is not a required dependency for `pylibraft`.
+The example below demonstrates computing the pairwise Euclidean distances between CuPy arrays. Note that CuPy is not a required dependency for `pylibraft`.
+
+```python
+import cupy as cp
+
+from pylibraft.distance import pairwise_distance
+
+n_samples = 5000
+n_features = 50
+
+in1 = cp.random.random_sample((n_samples, n_features), dtype=cp.float32)
+in2 = cp.random.random_sample((n_samples, n_features), dtype=cp.float32)
+
+output = pairwise_distance(in1, in2, metric="euclidean")
+```
+
+The `output` array supports `__cuda_array_interface__` so it is interoperable with other libraries like CuPy and PyTorch that also support it. 
+
+Below is an example of converting the output `pylibraft.device_ndarray` to a CuPy array:
+```python
+cupy_array = cp.asarray(output)
+```
+
+And converting to a PyTorch tensor:
+```python
+import torch
+
+torch_tensor = torch.as_tensor(output, device='cuda')
+```
+
+`pylibraft` also supports writing to a pre-allocated output array so any `__cuda_array_interface__` supported array can be written to in-place:
 
 ```python
 import cupy as cp
@@ -95,9 +157,10 @@ in1 = cp.random.random_sample((n_samples, n_features), dtype=cp.float32)
 in2 = cp.random.random_sample((n_samples, n_features), dtype=cp.float32)
 output = cp.empty((n_samples, n_samples), dtype=cp.float32)
 
-pairwise_distance(in1, in2, output, metric="euclidean")
+pairwise_distance(in1, in2, out=output, metric="euclidean")
 ```
 
+
 ## Installing
 
 RAFT itself can be installed through conda, [Cmake Package Manager (CPM)](https://github.com/cpm-cmake/CPM.cmake), or by building the repository from source. Please refer to the [build instructions](docs/source/build.md) for more a comprehensive guide on building RAFT and using it in downstream projects.

From c4f9c7f9fb2acf861bc25f33fc1de0695709ec5d Mon Sep 17 00:00:00 2001
From: "Corey J. Nolet" <cjnolet@gmail.com>
Date: Mon, 14 Nov 2022 14:02:48 -0500
Subject: [PATCH 11/14] Adding example of converting device_ndarray to torch
 tensor

---
 .../pylibraft/pylibraft/common/device_ndarray.py   | 14 ++++++++++++--
 1 file changed, 12 insertions(+), 2 deletions(-)

diff --git a/python/pylibraft/pylibraft/common/device_ndarray.py b/python/pylibraft/pylibraft/common/device_ndarray.py
index 6eee4ef04d..de821544d4 100644
--- a/python/pylibraft/pylibraft/common/device_ndarray.py
+++ b/python/pylibraft/pylibraft/common/device_ndarray.py
@@ -36,9 +36,10 @@ def __init__(self, np_ndarray):
         --------
         The device_ndarray is __cuda_array_interface__ compliant so it is
         interoperable with other libraries that also support it, such as
-        CuPy and PyTorch. The following usage example demonstrates
-        converting a pylibraft.device_ndarray to a cupy.ndarray:
+        CuPy and PyTorch.
 
+        The following usage example demonstrates
+        converting a pylibraft.device_ndarray to a cupy.ndarray:
         .. code-block:: python
 
             import cupy as cp
@@ -46,6 +47,15 @@ def __init__(self, np_ndarray):
 
             raft_array = device_ndarray.empty((100, 50))
             cupy_array = cp.asarray(raft_array)
+
+        And the converting pylibraft.device_ndarray to a PyTorch tensor:
+        .. code-block:: python
+
+            import torch
+            from pylibraft import device_ndarray
+
+            raft_array = device_ndarray.empty((100, 50))
+            torch_tensor = torch.as_tensor(raft_array, device='cuda')
         """
         self.ndarray_ = np_ndarray
         order = "C" if self.c_contiguous else "F"

From ec3cc6e8e450029b8837e8fd0ebf80c91429624e Mon Sep 17 00:00:00 2001
From: "Corey J. Nolet" <cjnolet@gmail.com>
Date: Mon, 14 Nov 2022 14:07:42 -0500
Subject: [PATCH 12/14] Adding to pairwise distance pydocs

---
 .../pylibraft/distance/pairwise_distance.pyx  | 33 +++++++++++++++++--
 1 file changed, 31 insertions(+), 2 deletions(-)

diff --git a/python/pylibraft/pylibraft/distance/pairwise_distance.pyx b/python/pylibraft/pylibraft/distance/pairwise_distance.pyx
index 9c1780ddc0..dc4bd982f9 100644
--- a/python/pylibraft/pylibraft/distance/pairwise_distance.pyx
+++ b/python/pylibraft/pylibraft/distance/pairwise_distance.pyx
@@ -123,6 +123,7 @@ def distance(X, Y, out=None, metric="euclidean", p=2.0, handle=None):
     Examples
     --------
 
+    To compute pairwise distances on cupy arrays:
     .. code-block:: python
 
         import cupy as cp
@@ -137,17 +138,45 @@ def distance(X, Y, out=None, metric="euclidean", p=2.0, handle=None):
                                       dtype=cp.float32)
         in2 = cp.random.random_sample((n_samples, n_features),
                                       dtype=cp.float32)
-        output = cp.empty((n_samples, n_samples), dtype=cp.float32)
 
         # A single RAFT handle can optionally be reused across
         # pylibraft functions.
         handle = Handle()
         ...
-        pairwise_distance(in1, in2, output, metric="euclidean", handle=handle)
+        output = pairwise_distance(in1, in2, metric="euclidean", handle=handle)
         ...
         # pylibraft functions are often asynchronous so the
         # handle needs to be explicitly synchronized
         handle.sync()
+
+   It's also possible to write to a pre-allocated output array:
+   .. code-block:: python
+
+       import cupy as cp
+
+       from pylibraft.common import Handle
+       from pylibraft.distance import pairwise_distance
+
+       n_samples = 5000
+       n_features = 50
+
+       in1 = cp.random.random_sample((n_samples, n_features),
+                                     dtype=cp.float32)
+       in2 = cp.random.random_sample((n_samples, n_features),
+                                     dtype=cp.float32)
+       output = cp.empty((n_samples, n_samples), dtype=cp.float32)
+
+       # A single RAFT handle can optionally be reused across
+       # pylibraft functions.
+       handle = Handle()
+       ...
+       pairwise_distance(in1, in2, out=output,
+                         metric="euclidean", handle=handle)
+       ...
+       # pylibraft functions are often asynchronous so the
+       # handle needs to be explicitly synchronized
+       handle.sync()
+
    """
 
     x_cai = X.__cuda_array_interface__

From f3d59f3de81d86e7582c824b45baf35c1f4cc402 Mon Sep 17 00:00:00 2001
From: "Corey J. Nolet" <cjnolet@gmail.com>
Date: Mon, 14 Nov 2022 14:19:00 -0500
Subject: [PATCH 13/14] Making in-place output optional for fused_l2_nn_argmin

---
 .../pylibraft/common/device_ndarray.py        | 13 ++---
 .../pylibraft/distance/fused_l2_nn.pyx        | 51 ++++++++++++++++---
 .../pylibraft/test/test_fused_l2_argmin.py    |  8 +--
 3 files changed, 57 insertions(+), 15 deletions(-)

diff --git a/python/pylibraft/pylibraft/common/device_ndarray.py b/python/pylibraft/pylibraft/common/device_ndarray.py
index de821544d4..eebbca2f06 100644
--- a/python/pylibraft/pylibraft/common/device_ndarray.py
+++ b/python/pylibraft/pylibraft/common/device_ndarray.py
@@ -20,13 +20,14 @@
 
 class device_ndarray:
     """
-    pylibraft.device_ndarray is meant to be a very lightweight
+    pylibraft.common.device_ndarray is meant to be a very lightweight
     __cuda_array_interface__ wrapper around a numpy.ndarray.
     """
 
     def __init__(self, np_ndarray):
         """
-        Construct a pylibraft.device_ndarray wrapper around a numpy.ndarray
+        Construct a pylibraft.common.device_ndarray wrapper around a
+        numpy.ndarray
 
         Parameters
         ----------
@@ -39,20 +40,20 @@ def __init__(self, np_ndarray):
         CuPy and PyTorch.
 
         The following usage example demonstrates
-        converting a pylibraft.device_ndarray to a cupy.ndarray:
+        converting a pylibraft.common.device_ndarray to a cupy.ndarray:
         .. code-block:: python
 
             import cupy as cp
-            from pylibraft import device_ndarray
+            from pylibraft.common import device_ndarray
 
             raft_array = device_ndarray.empty((100, 50))
             cupy_array = cp.asarray(raft_array)
 
-        And the converting pylibraft.device_ndarray to a PyTorch tensor:
+        And the converting pylibraft.common.device_ndarray to a PyTorch tensor:
         .. code-block:: python
 
             import torch
-            from pylibraft import device_ndarray
+            from pylibraft.common import device_ndarray
 
             raft_array = device_ndarray.empty((100, 50))
             torch_tensor = torch.as_tensor(raft_array, device='cuda')
diff --git a/python/pylibraft/pylibraft/distance/fused_l2_nn.pyx b/python/pylibraft/pylibraft/distance/fused_l2_nn.pyx
index 73cd60058f..9597e3906e 100644
--- a/python/pylibraft/pylibraft/distance/fused_l2_nn.pyx
+++ b/python/pylibraft/pylibraft/distance/fused_l2_nn.pyx
@@ -26,7 +26,7 @@ from libcpp cimport bool
 
 from .distance_type cimport DistanceType
 
-from pylibraft.common import Handle
+from pylibraft.common import Handle, device_ndarray
 from pylibraft.common.handle import auto_sync_handle
 from pylibraft.common.handle cimport handle_t
 
@@ -62,7 +62,7 @@ cdef extern from "raft_distance/fused_l2_min_arg.hpp" \
 
 
 @auto_sync_handle
-def fused_l2_nn_argmin(X, Y, output, sqrt=True, handle=None):
+def fused_l2_nn_argmin(X, Y, out=None, sqrt=True, handle=None):
     """
     Compute the 1-nearest neighbors between X and Y using the L2 distance
 
@@ -77,6 +77,35 @@ def fused_l2_nn_argmin(X, Y, output, sqrt=True, handle=None):
     Examples
     --------
 
+    To compute the 1-nearest neighbors argmin:
+    .. code-block:: python
+
+        import cupy as cp
+
+        from pylibraft.common import Handle
+        from pylibraft.distance import fused_l2_nn_argmin
+
+        n_samples = 5000
+        n_clusters = 5
+        n_features = 50
+
+        in1 = cp.random.random_sample((n_samples, n_features),
+                                      dtype=cp.float32)
+        in2 = cp.random.random_sample((n_clusters, n_features),
+                                      dtype=cp.float32)
+
+        # A single RAFT handle can optionally be reused across
+        # pylibraft functions.
+        handle = Handle()
+        ...
+        output = fused_l2_nn_argmin(in1, in2, output, handle=handle)
+        ...
+        # pylibraft functions are often asynchronous so the
+        # handle needs to be explicitly synchronized
+        handle.sync()
+
+    The output can also be computed in-place on a preallocated
+    array:
     .. code-block:: python
 
         import cupy as cp
@@ -98,20 +127,30 @@ def fused_l2_nn_argmin(X, Y, output, sqrt=True, handle=None):
         # pylibraft functions.
         handle = Handle()
         ...
-        fused_l2_nn_argmin(in1, in2, output, handle=handle)
+        fused_l2_nn_argmin(in1, in2, out=output, handle=handle)
         ...
         # pylibraft functions are often asynchronous so the
         # handle needs to be explicitly synchronized
         handle.sync()
+
    """
 
     x_cai = X.__cuda_array_interface__
     y_cai = Y.__cuda_array_interface__
-    output_cai = output.__cuda_array_interface__
+
+    x_dt = np.dtype(x_cai["typestr"])
+    y_dt = np.dtype(y_cai["typestr"])
 
     m = x_cai["shape"][0]
     n = y_cai["shape"][0]
 
+    if out is None:
+        output = device_ndarray.empty((m,), dtype="int32")
+    else:
+        output = out
+
+    output_cai = output.__cuda_array_interface__
+
     x_k = x_cai["shape"][1]
     y_k = y_cai["shape"][1]
 
@@ -127,8 +166,6 @@ def fused_l2_nn_argmin(X, Y, output, sqrt=True, handle=None):
     handle = handle if handle is not None else Handle()
     cdef handle_t *h = <handle_t*><size_t>handle.getHandle()
 
-    x_dt = np.dtype(x_cai["typestr"])
-    y_dt = np.dtype(y_cai["typestr"])
     d_dt = np.dtype(output_cai["typestr"])
 
     x_c_contiguous = is_c_cont(x_cai, x_dt)
@@ -162,3 +199,5 @@ def fused_l2_nn_argmin(X, Y, output, sqrt=True, handle=None):
                             <bool>sqrt)
     else:
         raise ValueError("dtype %s not supported" % x_dt)
+
+    return output
diff --git a/python/pylibraft/pylibraft/test/test_fused_l2_argmin.py b/python/pylibraft/pylibraft/test/test_fused_l2_argmin.py
index cc54389a31..b05ad3d530 100644
--- a/python/pylibraft/pylibraft/test/test_fused_l2_argmin.py
+++ b/python/pylibraft/pylibraft/test/test_fused_l2_argmin.py
@@ -21,11 +21,12 @@
 from pylibraft.distance import fused_l2_nn_argmin
 
 
+@pytest.mark.parametrize("inplace", [True, False])
 @pytest.mark.parametrize("n_rows", [10, 100])
 @pytest.mark.parametrize("n_clusters", [5, 10])
 @pytest.mark.parametrize("n_cols", [3, 5])
 @pytest.mark.parametrize("dtype", [np.float32, np.float64])
-def test_fused_l2_nn_minarg(n_rows, n_cols, n_clusters, dtype):
+def test_fused_l2_nn_minarg(n_rows, n_cols, n_clusters, dtype, inplace):
     input1 = np.random.random_sample((n_rows, n_cols))
     input1 = np.asarray(input1, order="C").astype(dtype)
 
@@ -39,13 +40,14 @@ def test_fused_l2_nn_minarg(n_rows, n_cols, n_clusters, dtype):
 
     input1_device = device_ndarray(input1)
     input2_device = device_ndarray(input2)
-    output_device = device_ndarray(output)
+    output_device = device_ndarray(output) if inplace else None
 
     handle = Handle()
-    fused_l2_nn_argmin(
+    ret_output = fused_l2_nn_argmin(
         input1_device, input2_device, output_device, True, handle=handle
     )
     handle.sync()
+    output_device = ret_output if not inplace else output_device
     actual = output_device.copy_to_host()
 
     assert np.allclose(expected, actual, rtol=1e-4)

From ec934030a506b40118578231312615cfc81b02fa Mon Sep 17 00:00:00 2001
From: "Corey J. Nolet" <cjnolet@gmail.com>
Date: Mon, 14 Nov 2022 14:28:56 -0500
Subject: [PATCH 14/14] Adding link to cuda array interface on README

---
 README.md | 2 +-
 1 file changed, 1 insertion(+), 1 deletion(-)

diff --git a/README.md b/README.md
index e0938638d1..04dd2ff16d 100755
--- a/README.md
+++ b/README.md
@@ -129,7 +129,7 @@ in2 = cp.random.random_sample((n_samples, n_features), dtype=cp.float32)
 output = pairwise_distance(in1, in2, metric="euclidean")
 ```
 
-The `output` array supports `__cuda_array_interface__` so it is interoperable with other libraries like CuPy and PyTorch that also support it. 
+The `output` array supports [__cuda_array_interface__](https://numba.pydata.org/numba-doc/dev/cuda/cuda_array_interface.html#cuda-array-interface-version-2) so it is interoperable with other libraries like CuPy, Numba, and PyTorch that also support it. 
 
 Below is an example of converting the output `pylibraft.device_ndarray` to a CuPy array:
 ```python