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

Code before:
        The decision function values or probability estimates.
        - If shape is (n_samples,), converts to (n_samples, 2) with (-x, x).
        - If shape is (n_samples, 1), converts to (n_samples, 2) with (-x, x).
        - If shape is (n_samples, n_classes), returns unchanged.
        - For probability estimates, returns `numpy.log(decision_values + eps)`.

    eps : float
        Small positive value added to avoid log(0).

    Returns
    -------
    logits : ndarray of shape (n_samples, n_classes)
    """
    xp, _, device_ = get_namespace_and_device(decision_values, xp=xp)
    decision_values = check_array(
        decision_values, dtype=[xp.float64, xp.float32], ensure_2d=False
    )
    if (decision_values.ndim == 2) and (decision_values.shape[1] > 1):
        # Check if it is the output of predict_proba
        entries_zero_to_one = xp.all((decision_values >= 0) & (decision_values <= 1))
        row_sums_to_one = xp.all(
            xpx.isclose(
                xp.sum(decision_values, axis=1),
                xp.asarray(1.0, device=device_, dtype=decision_values.dtype),
            )
        )

        if entries_zero_to_one and row_sums_to_one:
            logits = xp.log(decision_values + eps)
        else:
            logits = decision_values

    elif (decision_values.ndim == 2) and (decision_values.shape[1] == 1):
        logits = xp.concat([-decision_values, decision_values], axis=1)

    elif decision_values.ndim == 1:
        decision_values = xp.reshape(decision_values, (-1, 1))
        logits = xp.concat([-decision_values, decision_values], axis=1)

    return logits


class _SigmoidCalibration(RegressorMixin, BaseEstimator):
    """Sigmoid regression model.

    Attributes
    ----------
    a_ : float
        The slope.

    b_ : float

Review comment (refactor):
Is this really necessary to explictly move `1.0` to the device and dtype of `decision_values`? I think that `xpx.isclose` should be able to accept a Python scalar directly and do the necesarray conversion automatically for convenience. If it is not the case, could you please open an upstream issue and add a `# TODO: simplify once upstream issue is fixed: https://github.com/data-apis/array-api-extra/issues/XYZ`.

Code after:
        The decision function values or probability estimates.
        - If shape is (n_samples,), converts to (n_samples, 2) with (-x, x).
        - If shape is (n_samples, 1), converts to (n_samples, 2) with (-x, x).
        - If shape is (n_samples, n_classes), returns unchanged.
        - For probability estimates, returns `numpy.log(decision_values + eps)`.

    eps : float
        Small positive value added to avoid log(0).

    Returns
    -------
    logits : ndarray of shape (n_samples, n_classes)
    """
    xp, _, device_ = get_namespace_and_device(decision_values, xp=xp)
    decision_values = check_array(
        decision_values, dtype=[xp.float64, xp.float32], ensure_2d=False
    )
    if (decision_values.ndim == 2) and (decision_values.shape[1] > 1):
        # Check if it is the output of predict_proba
        entries_zero_to_one = xp.all((decision_values >= 0) & (decision_values <= 1))
        # TODO: simplify once upstream issue is addressed
        # https://github.com/data-apis/array-api-extra/issues/478
        row_sums_to_one = xp.all(
            xpx.isclose(
                xp.sum(decision_values, axis=1),
                xp.asarray(1.0, device=device_, dtype=decision_values.dtype),
            )
        )

        if entries_zero_to_one and row_sums_to_one:
            logits = xp.log(decision_values + eps)
        else:
            logits = decision_values

    elif (decision_values.ndim == 2) and (decision_values.shape[1] == 1):
        logits = xp.concat([-decision_values, decision_values], axis=1)

    elif decision_values.ndim == 1:
        decision_values = xp.reshape(decision_values, (-1, 1))
        logits = xp.concat([-decision_values, decision_values], axis=1)

    return logits


class _SigmoidCalibration(RegressorMixin, BaseEstimator):
    """Sigmoid regression model.

    Attributes
    ----------
    a_ : float
        The slope.