diff --git a/python/cuml/datasets/blobs.py b/python/cuml/datasets/blobs.py index 44e28671d3..ca3ebf415f 100644 --- a/python/cuml/datasets/blobs.py +++ b/python/cuml/datasets/blobs.py @@ -65,7 +65,7 @@ def _get_centers(rs, centers, center_box, n_samples, n_features, dtype): return centers, n_centers -@cuml.internals.api_return_any() +@cuml.internals.api_return_generic() def make_blobs(n_samples=100, n_features=2, centers=None, cluster_std=1.0, center_box=(-10.0, 10.0), shuffle=True, random_state=None, return_centers=False, order='F', dtype='float32'): @@ -133,6 +133,12 @@ def make_blobs(n_samples=100, n_features=2, centers=None, cluster_std=1.0, -------- make_classification: a more intricate variant """ + + # Set the default output type to "cupy". This will be ignored if the user + # has set `cuml.global_output_type`. Only necessary for array generation + # methods that do not take an array as input + cuml.internals.set_api_output_type("cupy") + generator = _create_rs_generator(random_state=random_state) centers, n_centers = _get_centers(generator, centers, center_box, diff --git a/python/cuml/test/test_make_blobs.py b/python/cuml/test/test_make_blobs.py index 08da7172c6..d212a3ae98 100644 --- a/python/cuml/test/test_make_blobs.py +++ b/python/cuml/test/test_make_blobs.py @@ -16,7 +16,9 @@ import cuml import pytest import cupy as cp - +import cudf +import numpy as np +import numba.cuda # Testing parameters for scalar parameter tests @@ -94,3 +96,30 @@ def test_make_blobs_scalar_parameters(dtype, n_samples, n_features, centers, elif centers <= n_samples: assert cp.unique(labels).shape == (centers,), \ "unexpected number of clusters" + + +test_output_types = { + None: cp.ndarray, # Default is cupy if None is used + 'numpy': np.ndarray, + 'cupy': cp.ndarray, + 'numba': numba.cuda.devicearray.DeviceNDArrayBase, + 'cudf': (cudf.DataFrame, cudf.Series) +} + + +@pytest.mark.parametrize("input_type", test_output_types.keys()) +def test_output_type(input_type: str): + + # Set the output type and ensure its respected by the function + with cuml.using_output_type(input_type): + X, y = cuml.make_blobs(n_samples=10, + centers=3, + n_features=2, + random_state=0) + + if (isinstance(test_output_types[input_type], tuple)): + assert (isinstance(X, test_output_types[input_type][0])) + assert (isinstance(y, test_output_types[input_type][1])) + else: + assert (isinstance(X, test_output_types[input_type])) + assert (isinstance(y, test_output_types[input_type]))