Repository: scikit-learn/scikit-learn
File: sklearn/utils/estimator_checks.py (Python)

Code before:


def check_array_api_input(
    name,
    estimator_orig,
    array_namespace,
    device=None,
    dtype_name="float64",
    check_values=False,
    check_sample_weight=False,
    expect_only_array_outputs=True,
):
    """Check that the estimator can work consistently with the Array API

    By default, this just checks that the types and shapes of the arrays are
    consistent with calling the same estimator with numpy arrays.

    When check_values is True, it also checks that calling the estimator on the
    array_api Array gives the same results as ndarrays.

    When check_sample_weight is True, dummy sample weights are passed to the
    fit call.

    When expect_only_array_outputs is False, the check is looser: in particular
    it accepts non-array outputs such as sparse data structures. This is
    useful to test that enabling array API dispatch does not change the 
    behavior of any estimator fed with NumPy inputs, even for estimators that
    do not support array API.
    """
    xp = _array_api_for_tests(array_namespace, device)

    X, y = make_classification(n_samples=30, n_features=10, random_state=42)
    X = X.astype(dtype_name, copy=False)

    X = _enforce_estimator_tags_X(estimator_orig, X)
    y = _enforce_estimator_tags_y(estimator_orig, y)

    est = clone(estimator_orig)
    set_random_state(est)

    X_xp = xp.asarray(X, device=device)
    y_xp = xp.asarray(y, device=device)
    fit_kwargs = {}
    fit_kwargs_xp = {}
    if check_sample_weight:
        fit_kwargs["sample_weight"] = np.ones(X.shape[0], dtype=X.dtype)
        fit_kwargs_xp["sample_weight"] = xp.asarray(
            fit_kwargs["sample_weight"], device=device
        )

    est.fit(X, y, **fit_kwargs)

Review comment (suggestion):
```suggestion
    useful to test that enabling array API dispatch does not change the
```

Code after:


def check_array_api_input(
    name,
    estimator_orig,
    array_namespace,
    device=None,
    dtype_name="float64",
    check_values=False,
    check_sample_weight=False,
    expect_only_array_outputs=True,
):
    """Check that the estimator can work consistently with the Array API

    By default, this just checks that the types and shapes of the arrays are
    consistent with calling the same estimator with numpy arrays.

    When check_values is True, it also checks that calling the estimator on the
    array_api Array gives the same results as ndarrays.

    When check_sample_weight is True, dummy sample weights are passed to the
    fit call.

    When expect_only_array_outputs is False, the check is looser: in particular
    it accepts non-array outputs such as sparse data structures. This is
    useful to test that enabling array API dispatch does not change the
    behavior of any estimator fed with NumPy inputs, even for estimators that
    do not support array API.
    """
    xp = _array_api_for_tests(array_namespace, device)

    X, y = make_classification(n_samples=30, n_features=10, random_state=42)
    X = X.astype(dtype_name, copy=False)

    X = _enforce_estimator_tags_X(estimator_orig, X)
    y = _enforce_estimator_tags_y(estimator_orig, y)

    est = clone(estimator_orig)
    set_random_state(est)

    X_xp = xp.asarray(X, device=device)
    y_xp = xp.asarray(y, device=device)
    fit_kwargs = {}
    fit_kwargs_xp = {}
    if check_sample_weight:
        fit_kwargs["sample_weight"] = np.ones(X.shape[0], dtype=X.dtype)
        fit_kwargs_xp["sample_weight"] = xp.asarray(
            fit_kwargs["sample_weight"], device=device
        )

    est.fit(X, y, **fit_kwargs)