Merge pull request #7671 from andfoy/add_kdtree

Add KDTree to cupyx.scipy.spatial
cupy · Feb 13, 2024 · aa6356f · aa6356f
2 parents e26bcae + 0a8a3be
commit aa6356f
Show file tree

Hide file tree

Showing 8 changed files with 1,912 additions and 4 deletions.
diff --git a/.pfnci/linux/tests/cuda-rapids.Dockerfile b/.pfnci/linux/tests/cuda-rapids.Dockerfile
@@ -1,9 +1,11 @@
-ARG BASE_IMAGE="rapidsai/rapidsai-core-dev:23.06-cuda11.8-devel-ubuntu22.04-py3.10"
+ARG BASE_IMAGE="rapidsai/base:23.12-cuda12.0-py3.10"
 FROM ${BASE_IMAGE}
 
+USER root
+
 RUN export DEBIAN_FRONTEND=noninteractive && \
     apt-get -qqy update && \
-    apt-get -qqy install ccache git curl
+    apt-get -qqy install build-essential cuda-toolkit-12-0 ccache git curl
 
+ENV CUDA_PATH "/usr/local/cuda"
 ENV PATH "/usr/lib/ccache:${PATH}"
-ENV DISABLE_JUPYTER true
diff --git a/.pfnci/linux/tests/cuda-rapids.sh b/.pfnci/linux/tests/cuda-rapids.sh
@@ -14,5 +14,6 @@ export NVCC="ccache nvcc"
 CXX=/usr/lib/ccache/g++ CUDA_PATH="/usr/local/cuda" pip install -v ".[test]"
 "$ACTIONS/unittest.sh" "not slow" \
     "cupyx_tests/scipy_tests/sparse_tests/csgraph_tests" \
-    "cupyx_tests/scipy_tests/spatial_tests"
+    "cupyx_tests/scipy_tests/spatial_tests/test_kdtree_pylibraft.py" \
+    "cupyx_tests/scipy_tests/spatial_tests/test_distance.py"
 "$ACTIONS/cleanup.sh"
diff --git a/cupyx/scipy/spatial/__init__.py b/cupyx/scipy/spatial/__init__.py
@@ -1,2 +1,3 @@
 from cupyx.scipy.spatial.distance import distance_matrix   # NOQA
 from cupyx.scipy.spatial._delaunay import Delaunay  # NOQA
+from cupyx.scipy.spatial._kdtree import KDTree   # NOQA