Repository: scikit-learn/scikit-learn
File: sklearn/tests/test_calibration.py (Python)

Code before:
    yield_namespace_device_dtype_combinations(),
    ids=_get_namespace_device_dtype_ids,
)
def test_temperature_scaling_array_api_compliance(array_namespace, device_, dtype_name):
    """Check temperature scaling calibration"""
    xp = _array_api_for_tests(array_namespace, device_)
    X, y = make_classification(
        n_samples=1000,
        n_features=10,
        n_informative=10,
        n_redundant=0,
        n_classes=5,
        n_clusters_per_class=1,
        class_sep=2.0,
        random_state=42,
    )
    X_train, X_cal, y_train, y_cal = train_test_split(X, y, random_state=42)
    X_train = X_train.astype(dtype_name)
    y_train = y_train.astype(dtype_name)

    X_cal = X_cal.astype(dtype_name)
    y_cal = y_cal.astype(dtype_name)
    X_cal_xp = xp.asarray(X_cal, device=device_)
    y_cal_xp = xp.asarray(y_cal, device=device_)

    clf = LogisticRegression(penalty=None, tol=1e-8, max_iter=200, random_state=0)
    clf.fit(X_train, y_train)

    with config_context(array_api_dispatch=True):
        cal_clf_xp = CalibratedClassifierCV(
            FrozenEstimator(clf), cv=3, method="temperature", ensemble=False
        ).fit(X_cal_xp, y_cal_xp)

    cal_clf_np = CalibratedClassifierCV(
        FrozenEstimator(clf), cv=3, method="temperature", ensemble=False
    ).fit(X_cal, y_cal)

    calibrator_np = cal_clf_np.calibrated_classifiers_[0].calibrators[0]
    calibrator_xp = cal_clf_xp.calibrated_classifiers_[0].calibrators[0]
    rtol = 1e-3 if device_ == "mps" else 1e-10
    assert np.isclose(calibrator_xp.beta_, calibrator_np.beta_, rtol=rtol)


Review comment (suggestion):
We should use a classifier that supports array API in this test. At the moment, we only have `LinearDiscriminantAnalysis` I believe.

Code after:


@pytest.mark.parametrize("ensemble", [False, True])
@pytest.mark.parametrize("use_sample_weight", [False, True])
@pytest.mark.parametrize(
    "array_namespace, device_, dtype_name",
    yield_namespace_device_dtype_combinations(),
    ids=_get_namespace_device_dtype_ids,
)
def test_temperature_scaling_array_api_compliance(
    ensemble, use_sample_weight, array_namespace, device_, dtype_name
):
    """Check that `CalibratedClassifierCV` with temperature scaling is compatible
    with the array API"""

    xp = _array_api_for_tests(array_namespace, device_)
    X, y = make_classification(
        n_samples=1000,
        n_features=10,
        n_informative=10,
        n_redundant=0,
        n_classes=5,
        n_clusters_per_class=1,
        class_sep=2.0,
        random_state=42,
    )
    X_train, X_cal, y_train, y_cal = train_test_split(X, y, random_state=42)

    X_train = X_train.astype(dtype_name)
    y_train = y_train.astype(dtype_name)
    X_train_xp = xp.asarray(X_train, device=device_)
    y_train_xp = xp.asarray(y_train, device=device_)

    X_cal = X_cal.astype(dtype_name)
    y_cal = y_cal.astype(dtype_name)
    X_cal_xp = xp.asarray(X_cal, device=device_)
    y_cal_xp = xp.asarray(y_cal, device=device_)

    if use_sample_weight:
        sample_weight = np.ones_like(y_cal)
        sample_weight[1::2] = 2
    else:
        sample_weight = None

    clf_np = LinearDiscriminantAnalysis()
    clf_np.fit(X_train, y_train)
    cal_clf_np = CalibratedClassifierCV(
        FrozenEstimator(clf_np), cv=3, method="temperature", ensemble=ensemble
    ).fit(X_cal, y_cal, sample_weight=sample_weight)

    calibrator_np = cal_clf_np.calibrated_classifiers_[0].calibrators[0]