Support coercing host arrays in mlfunc/ReflectedAttr

python/cuml/cuml/__init__.py

@@ -1,5 +1,5 @@
 #
-# SPDX-FileCopyrightText: Copyright (c) 2022-2026, NVIDIA CORPORATION.
+# SPDX-FileCopyrightText: Copyright (c) 2022-2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved.
 # SPDX-License-Identifier: Apache-2.0
 #
 
@@ -85,7 +85,9 @@ def _setup_cupy():
     # Enable rmm_cupy_allocator
     cp.cuda.set_allocator(rmm_cupy_allocator)
 
-    # XXX: workaround for https://github.com/cupy/cupy/issues/10084
+    # TODO: this is a workaround for https://github.com/cupy/cupy/issues/10084
+    # It can be conditionally done once the cupy fix is out (see
+    # https://github.com/rapidsai/cuml/issues/8364).
     copyreg.dispatch_table[cp.ndarray] = lambda x: (
         cp.array,
         (x.get(order="A"),),

python/cuml/cuml/fil/compat.py

@@ -1,5 +1,5 @@
 #
-# SPDX-FileCopyrightText: Copyright (c) 2022-2026, NVIDIA CORPORATION.
+# SPDX-FileCopyrightText: Copyright (c) 2022-2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved.
 # SPDX-License-Identifier: Apache-2.0
 #
 
@@ -15,12 +15,11 @@
 import nvforest
 import treelite
 
-from cuml.internals.array import CumlArray
 from cuml.internals.base import Base, get_handle
 from cuml.internals.device_type import DeviceType
 from cuml.internals.global_settings import GlobalSettings
 from cuml.internals.mixins import CMajorInputTagMixin
-from cuml.internals.outputs import reflect
+from cuml.internals.outputs import mlfunc
 from cuml.internals.validation import check_array
 
 
@@ -337,14 +336,14 @@ def get_dtype(self):
             raise RuntimeError("ForestInference not yet loaded")
         return self.model.forest.get_dtype()
 
-    @reflect
+    @mlfunc(preserve_index=True)
     def predict_proba(
         self,
         X,
         *,
         preds=None,
         chunk_size=None,
-    ) -> CumlArray:
+    ):
         if self.model is None:
             raise RuntimeError("ForestInference not yet loaded")
         if preds is not None:
@@ -358,46 +357,43 @@ def predict_proba(
                 nvforest.CPUForestInferenceClassifier,
             ),
         ):
-            X, index = check_array(
+            X = check_array(
                 X,
                 dtype=self.get_dtype(),
                 order="C",
                 mem_type="device"
                 if _is_nvforest_model_on_device(self.model)
                 else "host",
-                return_index=True,
                 ensure_all_finite=self.ensure_all_finite,
                 input_name="X",
             )
             out = self.model.predict_proba(X, chunk_size=chunk_size)
             mem_type = GlobalSettings().fil_memory_type.name
-            out = cp.asarray(out) if mem_type == "device" else cp.asnumpy(out)
-            return CumlArray(out, index=index)
+            return cp.asarray(out) if mem_type == "device" else cp.asnumpy(out)
         raise RuntimeError("Must be a classifier to run predict_proba()")
 
-    @reflect
+    @mlfunc(preserve_index=True)
     def predict(
         self,
         X,
         *,
         preds=None,
         chunk_size=None,
         threshold=None,
-    ) -> CumlArray:
+    ):
         if self.model is None:
             raise RuntimeError("ForestInference not yet loaded")
         if preds is not None:
             raise NotImplementedError(
                 "Setting preds argument is no longer supported"
             )
-        X, index = check_array(
+        X = check_array(
             X,
             dtype=self.get_dtype(),
             order="C",
             mem_type="device"
             if _is_nvforest_model_on_device(self.model)
             else "host",
-            return_index=True,
             ensure_all_finite=self.ensure_all_finite,
             input_name="X",
         )
@@ -424,56 +420,51 @@ def predict(
                 f"Unrecognized type for self.model: {type(self.model)}"
             )
         mem_type = GlobalSettings().fil_memory_type.name
-        out = cp.asarray(out) if mem_type == "device" else cp.asnumpy(out)
-        return CumlArray(out, index=index)
+        return cp.asarray(out) if mem_type == "device" else cp.asnumpy(out)
 
-    @reflect
+    @mlfunc(preserve_index=True)
     def predict_per_tree(self, X, *, preds=None, chunk_size=None):
         if self.model is None:
             raise RuntimeError("ForestInference not yet loaded")
         if preds is not None:
             raise NotImplementedError(
                 "Setting preds argument is no longer supported"
             )
-        X, index = check_array(
+        X = check_array(
             X,
             dtype=self.get_dtype(),
             order="C",
             mem_type="device"
             if _is_nvforest_model_on_device(self.model)
             else "host",
-            return_index=True,
             ensure_all_finite=self.ensure_all_finite,
             input_name="X",
         )
         out = self.model.predict_per_tree(X, chunk_size=chunk_size)
         mem_type = GlobalSettings().fil_memory_type.name
-        out = cp.asarray(out) if mem_type == "device" else cp.asnumpy(out)
-        return CumlArray(out, index=index)
+        return cp.asarray(out) if mem_type == "device" else cp.asnumpy(out)
 
-    @reflect
+    @mlfunc(preserve_index=True)
     def apply(self, X, *, preds=None, chunk_size=None):
         if self.model is None:
             raise RuntimeError("ForestInference not yet loaded")
         if preds is not None:
             raise NotImplementedError(
                 "Setting preds argument is no longer supported"
             )
-        X, index = check_array(
+        X = check_array(
             X,
             dtype=self.get_dtype(),
             order="C",
             mem_type="device"
             if _is_nvforest_model_on_device(self.model)
             else "host",
-            return_index=True,
             ensure_all_finite=self.ensure_all_finite,
             input_name="X",
         )
         out = self.model.apply(X, chunk_size=chunk_size)
         mem_type = GlobalSettings().fil_memory_type.name
-        out = cp.asarray(out) if mem_type == "device" else cp.asnumpy(out)
-        return CumlArray(out, index=index)
+        return cp.asarray(out) if mem_type == "device" else cp.asnumpy(out)
 
     def optimize(
         self,

python/cuml/cuml/internals/outputs.py

@@ -1,5 +1,5 @@
 #
-# SPDX-FileCopyrightText: Copyright (c) 2020-2026, NVIDIA CORPORATION.
+# SPDX-FileCopyrightText: Copyright (c) 2020-2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved.
 # SPDX-License-Identifier: Apache-2.0
 #
 import contextlib
@@ -431,6 +431,7 @@ def to_output(self, output_type=None, index=None):
                 if index is not None:
                     out.index = index
 
+        # At this point `out` is either a cupy or cudf object
         if output_type is None:
             # cupy when possible, cudf otherwise
             return out
@@ -467,12 +468,13 @@ def convert_arrays(obj, output_type="cupy", index=None, _legacy=False):
     Parameters
     ----------
     obj : object
-        The object to convert. Any cupy arrays (dense or sparse) or
-        cuml-specific output types (`ClassLabels`, `ArrayIndexPair`) will be
-        converted to the specified `output_type`. Some builtin collections
-        (dict, list, tuple) are traversed recursively to find array-likes.
-        Other array-likes (numpy, pandas, ...) will error as unsupported.
-        Any other type is passed through unchanged.
+        The object to convert. Any cupy arrays, numpy arrays, cupyx sparse
+        matrices, scipy sparse matrices, or cuml-specific output types
+        (`ClassLabels`, `ArrayIndexPair`) will be converted to the specified
+        `output_type`. Some builtin collections (dict, list, tuple) are
+        traversed recursively to find array-likes. Other array-likes (pandas,
+        ...) will error as unsupported. Any other type is passed through
+        unchanged.
     output_type : {'cupy', 'numpy', 'cudf', 'pandas', 'numba'}
         The output type to convert to.
     index : pandas.Index, cudf.Index, or None, default=None
@@ -501,7 +503,23 @@ def convert_arrays(obj, output_type="cupy", index=None, _legacy=False):
     if isinstance(obj, ClassLabels):
         return obj.to_output(output_type, index=index)
 
-    elif isinstance(obj, cp.ndarray):
+    if isinstance(obj, np.ndarray):
+        if output_type == "numpy":
+            return obj
+        elif output_type == "pandas":
+            if hasattr(index, "to_pandas"):
+                index = index.to_pandas()
+            if obj.ndim == 2:
+                if obj.shape[1] == 1:
+                    return pd.Series(obj.flatten(), index=index)
+                return pd.DataFrame(obj, index=index)
+            return pd.Series(obj, index=index)
+        else:
+            # Other output types use device memory, coerce to cupy and take
+            # cupy code path.
+            obj = cp.asarray(obj)
+
+    if isinstance(obj, cp.ndarray):
         if output_type == "numpy":
             return obj.get(order="A")
         elif output_type in (
@@ -563,14 +581,25 @@ def convert_arrays(obj, output_type="cupy", index=None, _legacy=False):
         else:
             return obj
 
+    elif sp.issparse(obj):
+        if output_type in ("numpy", "pandas"):
+            return obj
+        elif obj.format == "csr":
+            return cp_sp.csr_matrix(obj)
+        elif obj.format == "csc":
+            return cp_sp.csc_matrix(obj)
+        else:
+            # Use coo for coo and all other formats
+            return cp_sp.coo_matrix(obj)
+
     elif isinstance(
-        obj,
-        (np.ndarray, cudf.Series, cudf.DataFrame, pd.Series, pd.DataFrame),
+        obj, (cudf.Series, cudf.DataFrame, pd.Series, pd.DataFrame)
     ):
         raise TypeError(
             f"Cannot return objects of type {type(obj).__name__} directly "
             f"from an `mlfunc`-decorated function. Please return a "
-            f"`cupy.ndarray`, `cupyx.scipy.sparse.spmatrix`, `ArrayIndexPair`, "
+            f"`cupy.ndarray`, `numpy.ndarray`, `cupyx.scipy.sparse.spmatrix`, "
+            f"`scipy.sparse.spmatrix`, `ArrayIndexPair`, "
             f"or `ClassLabels` instead."
         )
 
@@ -604,23 +633,22 @@ def __init__(self, value):
 
         def _requires_reflection(self, obj) -> bool:
             """Check if `obj` requires reflection."""
-            if isinstance(obj, (cp.ndarray, ArrayIndexPair)):
+            if isinstance(obj, (cp.ndarray, np.ndarray, ArrayIndexPair)):
                 return True
-            elif cp_sp.issparse(obj):
+            elif cp_sp.issparse(obj) or sp.issparse(obj):
                 return True
             elif isinstance(
                 obj,
                 (
-                    np.ndarray,
                     cudf.Series,
                     cudf.DataFrame,
                     pd.Series,
                     pd.DataFrame,
                 ),
             ):
                 raise TypeError(
-                    "Array-like types other than cupy, cupyx.scipy.sparse, or "
-                    "`ArrayIndexPair` are not supported."
+                    "Array-like types other than cupy, cupyx.scipy.sparse, "
+                    "numpy, scipy.sparse, or `ArrayIndexPair` are not supported."
                 )
             elif isinstance(obj, (list, tuple)):
                 return any(self._requires_reflection(v) for v in obj)

python/cuml/tests/test_reflection.py

@@ -1,4 +1,4 @@
-# SPDX-FileCopyrightText: Copyright (c) 2020-2026, NVIDIA CORPORATION.
+# SPDX-FileCopyrightText: Copyright (c) 2020-2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved.
 # SPDX-License-Identifier: Apache-2.0
 import pickle
 
@@ -331,37 +331,54 @@ def test_global_input_with_estimator_output_type():
     assert_output_type(model.components_, "pandas")
 
 
+@pytest.mark.parametrize("input_type", ["numpy", "cupy"])
 @pytest.mark.parametrize("output_type", ["numpy", "cupy"])
 @pytest.mark.parametrize("order", ["C", "F"])
-def test_convert_arrays_dense_array(output_type, order):
-    X = cp.asarray(rand_array("cupy"), order=order)
+def test_convert_arrays_dense_array(input_type, output_type, order):
+    if input_type == "cupy":
+        X = cp.asarray(rand_array("cupy"), order=order)
+    else:
+        X = np.asarray(rand_array("numpy"), order=order)
 
     out = convert_arrays(X, output_type)
     assert_output_type(out, output_type)
     np.testing.assert_array_equal(cp.asnumpy(X), cp.asnumpy(out))
     assert out.flags.c_contiguous if order == "C" else out.flags.f_contiguous
 
 
+@pytest.mark.parametrize("input_type", ["scipy", "cupyx"])
+@pytest.mark.parametrize("format", ["coo", "csc", "csr"])
 @pytest.mark.parametrize("output_type", OUTPUT_TYPES)
-def test_convert_arrays_sparse_array(output_type):
-    X = cupyx.scipy.sparse.random(5, 5, random_state=42, density=0.5)
+def test_convert_arrays_sparse_array(input_type, output_type, format):
+    if input_type == "cupyx":
+        X = cupyx.scipy.sparse.random(
+            5, 5, random_state=42, density=0.5, format=format
+        )
+    else:
+        X = scipy.sparse.random(
+            5, 5, random_state=42, density=0.5, format=format
+        )
+
     out = convert_arrays(X, output_type)
 
     if output_type in ["cupy", "cudf", "numba"]:
         assert cupyx.scipy.sparse.issparse(out)
     else:
         assert scipy.sparse.issparse(out)
 
+    assert out.format == format
+
     np.testing.assert_array_equal(
         cp.asnumpy(X.todense()),
         cp.asnumpy(out.todense()),
     )
 
 
 @pytest.mark.parametrize("kind", ["dataframe", "series"])
+@pytest.mark.parametrize("input_type", ["cupy", "numpy"])
 @pytest.mark.parametrize("output_type", ["pandas", "cudf"])
-def test_convert_arrays_dataframe(kind, output_type):
-    arr = rand_array("cupy", shape=((8, 4) if kind == "dataframe" else 8))
+def test_convert_arrays_dataframe(kind, input_type, output_type):
+    arr = rand_array(input_type, shape=((8, 4) if kind == "dataframe" else 8))
     res = convert_arrays(arr, output_type)
     assert_output_type(res, output_type)
 
@@ -376,9 +393,12 @@ def test_convert_arrays_dataframe(kind, output_type):
 @pytest.mark.parametrize("xdf", [pd, cudf])
 @pytest.mark.parametrize("kind", ["dataframe", "series"])
 @pytest.mark.parametrize("use_pair", [False, True])
+@pytest.mark.parametrize("input_type", ["cupy", "numpy"])
 @pytest.mark.parametrize("output_type", ["pandas", "cudf"])
-def test_convert_arrays_dataframe_with_index(xdf, kind, use_pair, output_type):
-    arr = rand_array("cupy", shape=((8, 4) if kind == "dataframe" else 8))
+def test_convert_arrays_dataframe_with_index(
+    xdf, kind, use_pair, input_type, output_type
+):
+    arr = rand_array(input_type, shape=((8, 4) if kind == "dataframe" else 8))
     index = xdf.Index(["a", "b", "c", "d", "e", "f", "g", "h"])
 
     if use_pair: