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

Code before:
                "The labels array needs to contain at least two "
                "labels, got {0}.".format(lb.classes_)
            )

    transformed_labels = lb.transform(y_true)

    if transformed_labels.shape[1] == 1:
        transformed_labels = np.append(
            1 - transformed_labels, transformed_labels, axis=1
        )

    # If y_prob is of single dimension, assume y_true to be binary
    # and then check.
    if y_prob.ndim == 1:
        y_prob = y_prob[:, xp.newaxis]
    if y_prob.shape[1] == 1:
        y_prob = xp.concat([1 - y_prob, y_prob], axis=1)

    eps = xp.finfo(y_prob.dtype).eps

    # Make sure y_prob is normalized
    y_prob_sum = xp.sum(y_prob, axis=1)

    if not xp.all(
        xpx.isclose(
            y_prob_sum,
            xp.asarray(1, dtype=y_prob_sum.dtype, device=device_),
            rtol=sqrt(eps),
        )
    ):
        warnings.warn(
            "The y_prob values do not sum to one. Make sure to pass probabilities.",
            UserWarning,
        )

    # Check if dimensions are consistent.
    transformed_labels = check_array(transformed_labels)
    if len(lb.classes_) != y_prob.shape[1]:
        if labels is None:
            raise ValueError(
                "y_true and y_prob contain different number of "
                "classes: {0} vs {1}. Please provide the true "
                "labels explicitly through the labels argument. "
                "Classes found in "
                "y_true: {2}".format(
                    transformed_labels.shape[1], y_prob.shape[1], lb.classes_
                )
            )
        else:
            raise ValueError(
                "The number of classes in labels is different "

Review comment (suggestion):
Updated this because we follow the namespace and device of `y_prob` as `y_true` can or cannot be on the `xp namespace and device`

Code after:

    labels : array-like, default=None
        If not provided, labels will be inferred from y_true. If `labels`
        is `None` and `y_prob` has shape `(n_samples,)` the labels are
        assumed to be binary and are inferred from `y_true`.

    Returns
    -------
    transformed_labels : array of shape (n_samples, n_classes)

    y_prob : array of shape (n_samples, n_classes)
    """
    xp, _, device_ = get_namespace_and_device(y_prob)

    y_prob = check_array(
        y_prob, ensure_2d=False, dtype=supported_float_dtypes(xp, device=device_)
    )

    if xp.max(y_prob) > 1:
        raise ValueError(f"y_prob contains values greater than 1: {xp.max(y_prob)}")
    if xp.min(y_prob) < 0:
        raise ValueError(f"y_prob contains values lower than 0: {xp.min(y_prob)}")

    check_consistent_length(y_prob, y_true, sample_weight)
    if sample_weight is not None:
        _check_sample_weight(sample_weight, y_prob, force_float_dtype=False)

    transformed_labels, lb_classes = _one_hot_encoding_multiclass_target(
        y_true=y_true, labels=labels, target_xp=xp, target_device=device_
    )

    # If y_prob is of single dimension, assume y_true to be binary
    # and then check.
    if y_prob.ndim == 1:
        y_prob = y_prob[:, xp.newaxis]
    if y_prob.shape[1] == 1:
        y_prob = xp.concat([1 - y_prob, y_prob], axis=1)

    eps = xp.finfo(y_prob.dtype).eps

    # Make sure y_prob is normalized
    y_prob_sum = xp.sum(y_prob, axis=1)

    if not xp.all(
        xpx.isclose(
            y_prob_sum,
            xp.asarray(1, dtype=y_prob_sum.dtype, device=device_),
            rtol=sqrt(eps),
        )
    ):
        warnings.warn(