From 08ce1d773878a9534910c70454e92f7c32ec5ffe Mon Sep 17 00:00:00 2001
From: Dante Gama Dessavre <danteg@nvidia.com>
Date: Tue, 10 Dec 2024 21:12:40 -0600
Subject: [PATCH] FIX missing imports and other small fixes

---
 python/cuml/cuml/internals/base.pyx            |  6 +++++-
 python/cuml/cuml/internals/input_utils.py      |  5 +++++
 .../experimental/accel/test_sparse_dispatch.py | 18 ++++++++++++------
 3 files changed, 22 insertions(+), 7 deletions(-)

diff --git a/python/cuml/cuml/internals/base.pyx b/python/cuml/cuml/internals/base.pyx
index 73fe080b0c..787aec0447 100644
--- a/python/cuml/cuml/internals/base.pyx
+++ b/python/cuml/cuml/internals/base.pyx
@@ -37,6 +37,7 @@ except ImportError:
 
 import cuml
 import cuml.common
+from cuml.common.sparse_utils import is_sparse
 import cuml.internals.logger as logger
 import cuml.internals
 import cuml.internals.input_utils
@@ -47,6 +48,7 @@ from cuml.internals.input_utils import (
     determine_array_type,
     input_to_cuml_array,
     input_to_host_array,
+    input_to_host_array_with_sparse_support,
     is_array_like
 )
 from cuml.internals.memory_utils import determine_array_memtype
@@ -676,7 +678,9 @@ class UniversalBase(Base):
 
     def args_to_cpu(self, *args, **kwargs):
         # put all the args on host
-        new_args = tuple(input_to_host_array_with_sparse_support(arg) for arg in args)
+        new_args = tuple(
+            input_to_host_array_with_sparse_support(arg) for arg in args
+        )
 
         # put all the kwargs on host
         new_kwargs = dict()
diff --git a/python/cuml/cuml/internals/input_utils.py b/python/cuml/cuml/internals/input_utils.py
index cda8ff0ed8..f06841e2f9 100644
--- a/python/cuml/cuml/internals/input_utils.py
+++ b/python/cuml/cuml/internals/input_utils.py
@@ -530,6 +530,11 @@ def input_to_host_array(
 
 
 def input_to_host_array_with_sparse_support(X):
+    try:
+        if scipy_sparse.isspmatrix(X):
+            return X
+    except UnavailableError:
+        pass
     _array_type, is_sparse = determine_array_type_full(X)
     if is_sparse:
         if _array_type == "cupy":
diff --git a/python/cuml/cuml/tests/experimental/accel/test_sparse_dispatch.py b/python/cuml/cuml/tests/experimental/accel/test_sparse_dispatch.py
index 1229a4af47..f49834ed9d 100644
--- a/python/cuml/cuml/tests/experimental/accel/test_sparse_dispatch.py
+++ b/python/cuml/cuml/tests/experimental/accel/test_sparse_dispatch.py
@@ -22,10 +22,17 @@
 from sklearn.decomposition import TruncatedSVD
 from sklearn.kernel_ridge import KernelRidge
 from sklearn.linear_model import (
-    LinearRegression, LogisticRegression, ElasticNet,
-    Ridge, Lasso
+    LinearRegression,
+    LogisticRegression,
+    ElasticNet,
+    Ridge,
+    Lasso,
+)
+from sklearn.neighbors import (
+    NearestNeighbors,
+    KNeighborsClassifier,
+    KNeighborsRegressor,
 )
-from sklearn.neighbors import NearestNeighbors, KNeighborsClassifier, KNeighborsRegressor
 from sklearn.base import is_classifier, is_regressor
 
 
@@ -44,15 +51,14 @@
 
 @pytest.mark.parametrize("estimator_name", list(estimators.keys()))
 def test_sparse_support(estimator_name):
-    X_sparse = csr_matrix([[0, 1],
-                           [1, 0]])
+    X_sparse = csr_matrix([[0, 1], [1, 0]])
     print(X_sparse.shape[0])
     y_class = np.array([0, 1])
     y_reg = np.array([0.0, 1.0])
     estimator = estimators[estimator_name]()
     # Fit or fit_transform depending on the estimator type
     if isinstance(estimator, (KMeans, DBSCAN, TruncatedSVD, NearestNeighbors)):
-        if hasattr(estimator, 'fit_transform'):
+        if hasattr(estimator, "fit_transform"):
             estimator.fit_transform(X_sparse)
         else:
             estimator.fit(X_sparse)