scikit-learn-intelex 2024.6.0__py310-none-win_amd64.whl → 2024.7.0__py310-none-win_amd64.whl

scikit_learn_intelex-2024.7.0.data/data/Lib/site-packages/sklearnex/spmd/neighbors/tests/test_neighbors_spmd.py

@@ -0,0 +1,288 @@
+# ==============================================================================
+# Copyright 2024 Intel 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 numpy.testing import assert_allclose
+
+from onedal.tests.utils._dataframes_support import (
+    _convert_to_dataframe,
+    get_dataframes_and_queues,
+)
+from sklearnex.tests._utils_spmd import (
+    _assert_unordered_allclose,
+    _generate_classification_data,
+    _generate_regression_data,
+    _get_local_tensor,
+    _mpi_libs_and_gpu_available,
+    _spmd_assert_allclose,
+)
+
+
+@pytest.mark.skipif(
+    not _mpi_libs_and_gpu_available,
+    reason="GPU device and MPI libs required for test",
+)
+@pytest.mark.parametrize(
+    "dataframe,queue",
+    get_dataframes_and_queues(dataframe_filter_="dpnp,dpctl", device_filter_="gpu"),
+)
+@pytest.mark.mpi
+def test_knncls_spmd_gold(dataframe, queue):
+    # Import spmd and batch algo
+    from sklearnex.neighbors import KNeighborsClassifier as KNeighborsClassifier_Batch
+    from sklearnex.spmd.neighbors import KNeighborsClassifier as KNeighborsClassifier_SPMD
+
+    # Create gold data and convert to dataframe
+    X_train = np.array(
+        [
+            [0.0, 0.0],
+            [0.0, 1.0],
+            [1.0, 0.0],
+            [0.0, 2.0],
+            [2.0, 0.0],
+            [0.9, 1.0],
+            [0.0, -1.0],
+            [-1.0, 0.0],
+            [-1.0, -1.0],
+        ]
+    )
+    # TODO: handle situations where not all classes are present on all ranks?
+    y_train = np.array([0, 1, 0, 1, 0, 1, 0, 1, 0])
+    X_test = np.array(
+        [
+            [1.0, -0.5],
+            [-5.0, 1.0],
+            [0.0, 1.0],
+            [10.0, -10.0],
+        ]
+    )
+
+    local_dpt_X_train = _convert_to_dataframe(
+        _get_local_tensor(X_train), sycl_queue=queue, target_df=dataframe
+    )
+    local_dpt_y_train = _convert_to_dataframe(
+        _get_local_tensor(y_train), sycl_queue=queue, target_df=dataframe
+    )
+    local_dpt_X_test = _convert_to_dataframe(
+        _get_local_tensor(X_test), sycl_queue=queue, target_df=dataframe
+    )
+
+    # Ensure predictions of batch algo match spmd
+    spmd_model = KNeighborsClassifier_SPMD(n_neighbors=1, algorithm="brute").fit(
+        local_dpt_X_train, local_dpt_y_train
+    )
+    batch_model = KNeighborsClassifier_Batch(n_neighbors=1, algorithm="brute").fit(
+        X_train, y_train
+    )
+    spmd_dists, spmd_indcs = spmd_model.kneighbors(local_dpt_X_test)
+    batch_dists, batch_indcs = batch_model.kneighbors(X_test)
+    spmd_result = spmd_model.predict(local_dpt_X_test)
+    batch_result = batch_model.predict(X_test)
+
+    _assert_unordered_allclose(spmd_indcs, batch_indcs, localize=True)
+    _assert_unordered_allclose(spmd_dists, batch_dists, localize=True)
+    _spmd_assert_allclose(spmd_result, batch_result)
+
+
+@pytest.mark.skipif(
+    not _mpi_libs_and_gpu_available,
+    reason="GPU device and MPI libs required for test",
+)
+@pytest.mark.parametrize("n_samples", [200, 10000])
+@pytest.mark.parametrize("n_features_and_classes", [(5, 2), (25, 2), (25, 10)])
+@pytest.mark.parametrize("n_neighbors", [1, 5, 20])
+@pytest.mark.parametrize("weights", ["uniform", "distance"])
+@pytest.mark.parametrize(
+    "dataframe,queue",
+    get_dataframes_and_queues(dataframe_filter_="dpnp,dpctl", device_filter_="gpu"),
+)
+@pytest.mark.parametrize("dtype", [np.float32, np.float64])
+@pytest.mark.mpi
+def test_knncls_spmd_synthetic(
+    n_samples,
+    n_features_and_classes,
+    n_neighbors,
+    weights,
+    dataframe,
+    queue,
+    dtype,
+    metric="euclidean",
+):
+    n_features, n_classes = n_features_and_classes
+    # Import spmd and batch algo
+    from sklearnex.neighbors import KNeighborsClassifier as KNeighborsClassifier_Batch
+    from sklearnex.spmd.neighbors import KNeighborsClassifier as KNeighborsClassifier_SPMD
+
+    # Generate data and convert to dataframe
+    X_train, X_test, y_train, _ = _generate_classification_data(
+        n_samples, n_features, n_classes, dtype=dtype
+    )
+
+    local_dpt_X_train = _convert_to_dataframe(
+        _get_local_tensor(X_train), sycl_queue=queue, target_df=dataframe
+    )
+    local_dpt_y_train = _convert_to_dataframe(
+        _get_local_tensor(y_train), sycl_queue=queue, target_df=dataframe
+    )
+    local_dpt_X_test = _convert_to_dataframe(
+        _get_local_tensor(X_test), sycl_queue=queue, target_df=dataframe
+    )
+
+    # Ensure predictions of batch algo match spmd
+    spmd_model = KNeighborsClassifier_SPMD(
+        n_neighbors=n_neighbors, weights=weights, metric=metric, algorithm="brute"
+    ).fit(local_dpt_X_train, local_dpt_y_train)
+    batch_model = KNeighborsClassifier_Batch(
+        n_neighbors=n_neighbors, weights=weights, metric=metric, algorithm="brute"
+    ).fit(X_train, y_train)
+    spmd_dists, spmd_indcs = spmd_model.kneighbors(local_dpt_X_test)
+    batch_dists, batch_indcs = batch_model.kneighbors(X_test)
+    spmd_result = spmd_model.predict(local_dpt_X_test)
+    batch_result = batch_model.predict(X_test)
+
+    tol = 1e-4
+    if dtype == np.float64:
+        _assert_unordered_allclose(spmd_indcs, batch_indcs, localize=True)
+        _assert_unordered_allclose(
+            spmd_dists, batch_dists, localize=True, rtol=tol, atol=tol
+        )
+    _spmd_assert_allclose(spmd_result, batch_result)
+
+
+@pytest.mark.skipif(
+    not _mpi_libs_and_gpu_available,
+    reason="GPU device and MPI libs required for test",
+)
+@pytest.mark.parametrize(
+    "dataframe,queue",
+    get_dataframes_and_queues(dataframe_filter_="dpnp,dpctl", device_filter_="gpu"),
+)
+@pytest.mark.mpi
+def test_knnreg_spmd_gold(dataframe, queue):
+    # Import spmd and batch algo
+    from sklearnex.neighbors import KNeighborsRegressor as KNeighborsRegressor_Batch
+    from sklearnex.spmd.neighbors import KNeighborsRegressor as KNeighborsRegressor_SPMD
+
+    # Create gold data and convert to dataframe
+    X_train = np.array(
+        [
+            [0.0, 0.0],
+            [0.0, 1.0],
+            [1.0, 0.0],
+            [0.0, 2.0],
+            [2.0, 0.0],
+            [1.0, 1.0],
+            [0.0, -1.0],
+            [-1.0, 0.0],
+            [-1.0, -1.0],
+        ]
+    )
+    y_train = np.array([3.0, 5.0, 4.0, 7.0, 5.0, 6.0, 1.0, 2.0, 0.0])
+    X_test = np.array(
+        [
+            [1.0, -0.5],
+            [-5.0, 1.0],
+            [0.0, 1.0],
+            [10.0, -10.0],
+        ]
+    )
+
+    local_dpt_X_train = _convert_to_dataframe(
+        _get_local_tensor(X_train), sycl_queue=queue, target_df=dataframe
+    )
+    local_dpt_y_train = _convert_to_dataframe(
+        _get_local_tensor(y_train), sycl_queue=queue, target_df=dataframe
+    )
+    local_dpt_X_test = _convert_to_dataframe(
+        _get_local_tensor(X_test), sycl_queue=queue, target_df=dataframe
+    )
+
+    # Ensure predictions of batch algo match spmd
+    spmd_model = KNeighborsRegressor_SPMD(n_neighbors=1, algorithm="brute").fit(
+        local_dpt_X_train, local_dpt_y_train
+    )
+    batch_model = KNeighborsRegressor_Batch(n_neighbors=1, algorithm="brute").fit(
+        X_train, y_train
+    )
+    spmd_dists, spmd_indcs = spmd_model.kneighbors(local_dpt_X_test)
+    batch_dists, batch_indcs = batch_model.kneighbors(X_test)
+    spmd_result = spmd_model.predict(local_dpt_X_test)
+    batch_result = batch_model.predict(X_test)
+
+    _assert_unordered_allclose(spmd_indcs, batch_indcs, localize=True)
+    _assert_unordered_allclose(spmd_dists, batch_dists, localize=True)
+    _spmd_assert_allclose(spmd_result, batch_result)
+
+
+@pytest.mark.skipif(
+    not _mpi_libs_and_gpu_available,
+    reason="GPU device and MPI libs required for test",
+)
+@pytest.mark.parametrize("n_samples", [200, 10000])
+@pytest.mark.parametrize("n_features", [5, 25])
+@pytest.mark.parametrize("n_neighbors", [1, 5, 20])
+@pytest.mark.parametrize("weights", ["uniform", "distance"])
+@pytest.mark.parametrize(
+    "metric", ["euclidean", "manhattan", "minkowski", "chebyshev", "cosine"]
+)
+@pytest.mark.parametrize(
+    "dataframe,queue",
+    get_dataframes_and_queues(dataframe_filter_="dpnp,dpctl", device_filter_="gpu"),
+)
+@pytest.mark.parametrize("dtype", [np.float32, np.float64])
+@pytest.mark.mpi
+def test_knnreg_spmd_synthetic(
+    n_samples, n_features, n_neighbors, weights, metric, dataframe, queue, dtype
+):
+    # Import spmd and batch algo
+    from sklearnex.neighbors import KNeighborsRegressor as KNeighborsRegressor_Batch
+    from sklearnex.spmd.neighbors import KNeighborsRegressor as KNeighborsRegressor_SPMD
+
+    # Generate data and convert to dataframe
+    X_train, X_test, y_train, _ = _generate_regression_data(
+        n_samples, n_features, dtype=dtype
+    )
+
+    local_dpt_X_train = _convert_to_dataframe(
+        _get_local_tensor(X_train), sycl_queue=queue, target_df=dataframe
+    )
+    local_dpt_y_train = _convert_to_dataframe(
+        _get_local_tensor(y_train), sycl_queue=queue, target_df=dataframe
+    )
+    local_dpt_X_test = _convert_to_dataframe(
+        _get_local_tensor(X_test), sycl_queue=queue, target_df=dataframe
+    )
+
+    # Ensure predictions of batch algo match spmd
+    spmd_model = KNeighborsRegressor_SPMD(
+        n_neighbors=n_neighbors, weights=weights, metric=metric, algorithm="brute"
+    ).fit(local_dpt_X_train, local_dpt_y_train)
+    batch_model = KNeighborsRegressor_Batch(
+        n_neighbors=n_neighbors, weights=weights, metric=metric, algorithm="brute"
+    ).fit(X_train, y_train)
+    spmd_dists, spmd_indcs = spmd_model.kneighbors(local_dpt_X_test)
+    batch_dists, batch_indcs = batch_model.kneighbors(X_test)
+    spmd_result = spmd_model.predict(local_dpt_X_test)
+    batch_result = batch_model.predict(X_test)
+
+    tol = 0.005 if dtype == np.float32 else 1e-4
+    if dtype == np.float64:
+        _assert_unordered_allclose(spmd_indcs, batch_indcs, localize=True)
+        _assert_unordered_allclose(
+            spmd_dists, batch_dists, localize=True, rtol=tol, atol=tol
+        )
+    _spmd_assert_allclose(spmd_result, batch_result, rtol=tol, atol=tol)

{scikit_learn_intelex-2024.6.0.data → scikit_learn_intelex-2024.7.0.data}/data/Lib/site-packages/sklearnex/svm/_common.py

@@ -14,6 +14,7 @@
 # limitations under the License.
 # ==============================================================================
 
+import warnings
 from abc import ABC
 from numbers import Number, Real
 
@@ -22,7 +23,6 @@ from scipy import sparse as sp
 from sklearn.base import BaseEstimator, ClassifierMixin
 from sklearn.calibration import CalibratedClassifierCV
 from sklearn.metrics import r2_score
-from sklearn.model_selection import StratifiedKFold
 from sklearn.preprocessing import LabelEncoder
 
 from daal4py.sklearn._utils import sklearn_check_version
@@ -230,6 +230,17 @@ class BaseSVC(BaseSVM):
         return recip_freq[le.transform(classes)]
 
     def _fit_proba(self, X, y, sample_weight=None, queue=None):
+        # TODO: rewrite this method when probabilities output is implemented in oneDAL
+
+        # LibSVM uses the random seed to control cross-validation for probability generation
+        # CalibratedClassifierCV with "prefit" does not use an RNG nor a seed. This may
+        # impact users without their knowledge, so display a warning.
+        if self.random_state is not None:
+            warnings.warn(
+                "random_state does not influence oneDAL SVM results",
+                RuntimeWarning,
+            )
+
         params = self.get_params()
         params["probability"] = False
         params["decision_function_shape"] = "ovr"
@@ -240,26 +251,13 @@ class BaseSVC(BaseSVM):
         cfg = get_config()
         cfg["target_offload"] = queue
         with config_context(**cfg):
-            try:
-                n_splits = 5
-                n_jobs = n_splits if queue is None or queue.sycl_device.is_cpu else 1
-                cv = StratifiedKFold(
-                    n_splits=n_splits, shuffle=True, random_state=self.random_state
-                )
-                self.clf_prob = CalibratedClassifierCV(
-                    clf_base,
-                    ensemble=False,
-                    cv=cv,
-                    method="sigmoid",
-                )
-                self.clf_prob.fit(X, y, sample_weight)
-
-            except ValueError:
-                clf_base = clf_base.fit(X, y, sample_weight)
-                self.clf_prob = CalibratedClassifierCV(
-                    clf_base, cv="prefit", method="sigmoid"
-                )
-                self.clf_prob.fit(X, y, sample_weight)
+            clf_base.fit(X, y)
+            self.clf_prob = CalibratedClassifierCV(
+                clf_base,
+                ensemble=False,
+                cv="prefit",
+                method="sigmoid",
+            ).fit(X, y)
 
     def _save_attributes(self):
         self.support_vectors_ = self._onedal_estimator.support_vectors_

{scikit_learn_intelex-2024.6.0.data → scikit_learn_intelex-2024.7.0.data}/data/Lib/site-packages/sklearnex/svm/tests/test_svm.py

@@ -25,12 +25,10 @@ from onedal.tests.utils._dataframes_support import (
 )
 
 
-# TODO:
-# investigate failure for `dpnp.ndarrays` and `dpctl.tensors` on `GPU`
-@pytest.mark.parametrize(
-    "dataframe,queue", get_dataframes_and_queues(device_filter_="cpu")
-)
+@pytest.mark.parametrize("dataframe,queue", get_dataframes_and_queues())
 def test_sklearnex_import_svc(dataframe, queue):
+    if queue and queue.sycl_device.is_gpu:
+        pytest.skip("SVC fit for the GPU sycl_queue is buggy.")
     from sklearnex.svm import SVC
 
     X = np.array([[-2, -1], [-1, -1], [-1, -2], [+1, +1], [+1, +2], [+2, +1]])
@@ -43,12 +41,10 @@ def test_sklearnex_import_svc(dataframe, queue):
     assert_allclose(_as_numpy(svc.support_), [1, 3])
 
 
-# TODO:
-# investigate failure for `dpnp.ndarrays` and `dpctl.tensors` on `GPU`
-@pytest.mark.parametrize(
-    "dataframe,queue", get_dataframes_and_queues(device_filter_="cpu")
-)
+@pytest.mark.parametrize("dataframe,queue", get_dataframes_and_queues())
 def test_sklearnex_import_nusvc(dataframe, queue):
+    if queue and queue.sycl_device.is_gpu:
+        pytest.skip("NuSVC fit for the GPU sycl_queue is buggy.")
     from sklearnex.svm import NuSVC
 
     X = np.array([[-2, -1], [-1, -1], [-1, -2], [+1, +1], [+1, +2], [+2, +1]])
@@ -63,12 +59,10 @@ def test_sklearnex_import_nusvc(dataframe, queue):
     assert_allclose(_as_numpy(svc.support_), [0, 1, 3, 4])
 
 
-# TODO:
-# investigate failure for `dpnp.ndarrays` and `dpctl.tensors` on `GPU`
-@pytest.mark.parametrize(
-    "dataframe,queue", get_dataframes_and_queues(device_filter_="cpu")
-)
+@pytest.mark.parametrize("dataframe,queue", get_dataframes_and_queues())
 def test_sklearnex_import_svr(dataframe, queue):
+    if queue and queue.sycl_device.is_gpu:
+        pytest.skip("SVR fit for the GPU sycl_queue is buggy.")
     from sklearnex.svm import SVR
 
     X = np.array([[-2, -1], [-1, -1], [-1, -2], [+1, +1], [+1, +2], [+2, +1]])
@@ -81,12 +75,10 @@ def test_sklearnex_import_svr(dataframe, queue):
     assert_allclose(_as_numpy(svc.support_), [1, 3])
 
 
-# TODO:
-# investigate failure for `dpnp.ndarrays` and `dpctl.tensors` on `GPU`
-@pytest.mark.parametrize(
-    "dataframe,queue", get_dataframes_and_queues(device_filter_="cpu")
-)
+@pytest.mark.parametrize("dataframe,queue", get_dataframes_and_queues())
 def test_sklearnex_import_nusvr(dataframe, queue):
+    if queue and queue.sycl_device.is_gpu:
+        pytest.skip("NuSVR fit for the GPU sycl_queue is buggy.")
     from sklearnex.svm import NuSVR
 
     X = np.array([[-2, -1], [-1, -1], [-1, -2], [+1, +1], [+1, +2], [+2, +1]])

{scikit_learn_intelex-2024.6.0.data → scikit_learn_intelex-2024.7.0.data}/data/Lib/site-packages/sklearnex/tests/_utils.py

@@ -15,7 +15,7 @@
 # ==============================================================================
 
 from functools import partial
-from inspect import isclass
+from inspect import getattr_static, isclass, signature
 
 import numpy as np
 from scipy import sparse as sp
@@ -44,6 +44,22 @@ from sklearnex.svm import SVC, NuSVC
 
 
 def _load_all_models(with_sklearnex=True, estimator=True):
+    """Convert sklearnex patch_map into a dictionary of estimators or functions
+
+    Parameters
+    ----------
+    with_sklearnex: bool (default=True)
+        Discover estimators and methods with sklearnex patching enabled (True)
+        or disabled (False) from the sklearnex patch_map
+
+    estimator: bool (default=True)
+        yield estimators (True) or functions (False)
+
+    Returns
+    -------
+    dict: {name:estimator}
+        estimator is a class or function from sklearn or sklearnex
+    """
     # insure that patch state is correct as dictated by patch_sklearn boolean
     # and return it to the previous state no matter what occurs.
     already_patched_map = sklearn_is_patched(return_map=True)
@@ -91,11 +107,17 @@ mixin_map = [
 
 
 class _sklearn_clone_dict(dict):
+    """Special dict type for returning state-free sklearn/sklearnex estimators
+    with the same parameters"""
 
     def __getitem__(self, key):
         return clone(super().__getitem__(key))
 
 
+# Special dictionary of sklearnex estimators which must be specifically tested, this
+# could be because of supported non-default parameters, blocked support via sklearn's
+# 'available_if' decorator, or not being a native sklearn estimator (i.e. those not in
+# the default PATCHED_MODELS dictionary)
 SPECIAL_INSTANCES = _sklearn_clone_dict(
     {
         str(i): i
@@ -112,38 +134,107 @@ SPECIAL_INSTANCES = _sklearn_clone_dict(
 )
 
 
-def gen_models_info(algorithms):
+def gen_models_info(algorithms, required_inputs=["X", "y"]):
+    """Generate estimator-attribute pairs for pytest test collection.
+
+    Parameters
+    ----------
+    algorithms : iterable (list, tuple, 1D array-like object)
+        Iterable of valid sklearnex estimators or keys from PATCHED_MODELS
+
+    required_inputs : list, tuple of strings or None
+        list of required args/kwargs for callable attribute (only non-private,
+        non-BaseEstimator attributes).  Only one must be present, None
+        signifies taking all non-private attribues, callable or not.
+
+    Returns
+    -------
+    list of 2-element tuples: (estimator, string)
+        Returns a list of valid methods or attributes without "fit"
+    """
     output = []
-    for i in algorithms:
+    for estimator in algorithms:
 
-        if i in PATCHED_MODELS:
-            est = PATCHED_MODELS[i]
-        elif isinstance(algorithms[i], BaseEstimator):
-            est = algorithms[i].__class__
+        if estimator in PATCHED_MODELS:
+            est = PATCHED_MODELS[estimator]
+        elif isinstance(algorithms[estimator], BaseEstimator):
+            est = algorithms[estimator].__class__
+        else:
+            raise KeyError(f"Unrecognized sklearnex estimator: {estimator}")
+
+        # remove BaseEstimator methods (get_params, set_params)
+        candidates = set(dir(est)) - set(dir(BaseEstimator))
+        # remove private methods
+        candidates = set([attr for attr in candidates if not attr.startswith("_")])
+        # required to enable other methods
+        candidates = candidates - {"fit"}
+
+        # allow only callable methods with any of the required inputs
+        if required_inputs:
+            methods = []
+            for attr in candidates:
+                attribute = getattr_static(est, attr)
+                if callable(attribute):
+                    params = signature(attribute).parameters
+                    if any([inp in params for inp in required_inputs]):
+                        methods += [attr]
         else:
-            raise KeyError(f"Unrecognized sklearnex estimator: {i}")
+            methods = candidates
 
-        methods = set()
-        candidates = set(
-            [i for i in dir(est) if not i.startswith("_") and not i.endswith("_")]
+        output += (
+            [(estimator, method) for method in methods]
+            if methods
+            else [(estimator, None)]
         )
 
-        for mixin, method, _ in mixin_map:
-            if issubclass(est, mixin):
-                methods |= candidates & set(method)
-
-        output += [[i, j] for j in methods] if methods else [[i, None]]
-
     # In the case that no methods are available, set method to None.
     # This will allow estimators without mixins to still test the fit
     # method in various tests.
     return output
 
 
-def gen_dataset_type(est):
+def call_method(estimator, method, X, y, **kwargs):
+    """Generalized interface to call most sklearn estimator methods
+
+    Parameters
+    ----------
+    estimator : sklearn or sklearnex estimator instance
+
+    method: string
+        Valid callable method to estimator
+
+    X: array-like
+        data
+
+    y: array-like (for 'score', 'partial-fit', and 'path')
+        X-dependent data
+
+    **kwargs: keyword dict
+        keyword arguments to estimator.method
+
+    Returns
+    -------
+    return value from estimator.method
+    """
+    # useful for repository wide testing
+    if method == "inverse_transform":
+        # PCA's inverse_transform takes (n_samples, n_components)
+        data = (
+            (X[:, : estimator.n_components_],)
+            if X.shape[1] != estimator.n_components_
+            else (X,)
+        )
+    elif method not in ["score", "partial_fit", "path"]:
+        data = (X,)
+    else:
+        data = (X, y)
+    return getattr(estimator, method)(*data, **kwargs)
+
+
+def _gen_dataset_type(est):
     # est should be an estimator or estimator class
     # dataset initialized to classification, but will be swapped
-    # for other types as necessary
+    # for other types as necessary. Private method.
     dataset = "classification"
     estimator = est.__class__ if isinstance(est, BaseEstimator) else est
 
@@ -167,7 +258,39 @@ def gen_dataset(
     target_df=None,
     dtype=None,
 ):
-    dataset_type = gen_dataset_type(est)
+    """Generate dataset for pytest testing.
+
+    Parameters
+    ----------
+    est : sklearn or sklearnex estimator class
+        Must inherit an sklearn Mixin or sklearn's BaseEstimator
+
+    dataset: dataset dict
+        Dictionary with keys "classification" and/or "regression"
+        Value must be a list of object which yield X, y array
+        objects when called, ideally using a lambda or
+        functools.partial.
+
+    sparse: bool (default False)
+        Convert X data to a scipy.sparse csr_matrix format.
+
+    queue: SYCL queue or None
+        Queue necessary for device offloading following the
+        SYCL 2020 standard, usually generated by dpctl.
+
+    target_df: string or None
+        dataframe type for returned dataset, as dictated by
+        onedal's _convert_to_dataframe.
+
+    dtype: numpy dtype or None
+       target datatype for returned datasets (see DTYPES).
+
+    Returns
+    -------
+    list of 2-element list X,y: (array-like, array-like)
+        list of datasets for analysis
+    """
+    dataset_type = _gen_dataset_type(est)
     output = []
     # load data
     flag = dtype is None