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

Code before:
    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]
    pred_np = cal_clf_np.predict(X_train)
    with config_context(array_api_dispatch=True):
        clf_xp = LinearDiscriminantAnalysis()
        clf_xp.fit(X_train_xp, y_train_xp)
        cal_clf_xp = CalibratedClassifierCV(
            FrozenEstimator(clf_xp), cv=3, method="temperature", ensemble=ensemble
        ).fit(X_cal_xp, y_cal_xp, sample_weight=sample_weight)

        calibrator_xp = cal_clf_xp.calibrated_classifiers_[0].calibrators[0]
        rtol = 1e-3 if dtype_name == "float32" else 1e-7
        assert get_namespace(calibrator_xp.beta_)[0].__name__ == xp.__name__
        assert device(calibrator_xp.beta_) == device(X_cal_xp)
        assert np.isclose(
            _convert_to_numpy(calibrator_xp.beta_, xp=xp),
            calibrator_np.beta_,
            rtol=rtol,
        )
        pred_xp = cal_clf_xp.predict(X_train_xp)
        assert_allclose(_convert_to_numpy(pred_xp, xp=xp), pred_np)


@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_with_str_y_estimator_not_prefit(
    ensemble, use_sample_weight, array_namespace, device_, dtype_name
):
    """Check that `CalibratedClassifierCV` with temperature scaling is compatible
    with the array API when `y` is an ndarray of strings and the estimator is not
    fit beforehand (i.e. it is fit within `CalibratedClassifierCV`).
    """

    # TODO: Also ensure that `CalibratedClassifierCV` works appropriately with
    #  the array API when `y` is an ndarray of strings and we fit

Review comment (bug):
To get a more informative message in case of failure.

```suggestion
        assert_allclose(
```

Code after:
    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]
    pred_np = cal_clf_np.predict(X_train)
    with config_context(array_api_dispatch=True):
        clf_xp = LinearDiscriminantAnalysis()
        clf_xp.fit(X_train_xp, y_train_xp)
        cal_clf_xp = CalibratedClassifierCV(
            FrozenEstimator(clf_xp), cv=3, method="temperature", ensemble=ensemble
        ).fit(X_cal_xp, y_cal_xp, sample_weight=sample_weight)

        calibrator_xp = cal_clf_xp.calibrated_classifiers_[0].calibrators[0]
        rtol = 1e-3 if dtype_name == "float32" else 1e-7
        assert get_namespace(calibrator_xp.beta_)[0].__name__ == xp.__name__
        assert calibrator_xp.beta_.dtype == X_cal_xp.dtype
        assert device(calibrator_xp.beta_) == device(X_cal_xp)
        assert_allclose(
            _convert_to_numpy(calibrator_xp.beta_, xp=xp),
            calibrator_np.beta_,
            rtol=rtol,
        )
        pred_xp = cal_clf_xp.predict(X_train_xp)
        assert_allclose(_convert_to_numpy(pred_xp, xp=xp), pred_np)


@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_with_str_y_estimator_not_prefit(
    ensemble, use_sample_weight, array_namespace, device_, dtype_name
):
    """Check that `CalibratedClassifierCV` with temperature scaling is compatible
    with the array API when `y` is an ndarray of strings and the estimator is not
    fit beforehand (i.e. it is fit within `CalibratedClassifierCV`).
    """

    # TODO: Also ensure that `CalibratedClassifierCV` works appropriately with