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

Code before:
                Y = xp.zeros((n_samples, 1), dtype=xp.int64)
                Y += neg_label
                return Y
        elif n_classes >= 3:
            y_type = "multiclass"

    sorted_class = xp.sort(classes)
    if y_type == "multilabel-indicator":
        y_n_classes = y.shape[1] if hasattr(y, "shape") else len(y[0])
        if classes.size != y_n_classes:
            raise ValueError(
                "classes {0} mismatch with the labels {1} found in the data".format(
                    classes, unique_labels(y)
                )
            )

    if y_type in ("binary", "multiclass"):
        y = column_or_1d(y)

        # pick out the known labels from y
        y_in_classes = xp.asarray(
            xp.any(y[:, None] == classes[None, :], axis=1), dtype=xp.bool
        )
        y_in_classes_int = xp.asarray(
            xp.any(y[:, None] == classes[None, :], axis=1), dtype=xp.int64
        )
        y_seen = y[y_in_classes]
        indices = xp.searchsorted(sorted_class, y_seen)
        indptr = xp.concat(
            (
                xp.asarray([0], device=device_),
                xp.cumulative_sum(y_in_classes_int, axis=0),
            )
        )
        data = xp.full_like(indices, pos_label)

        if sparse_output:
            Y = sp.csr_matrix((data, indices, indptr), shape=(n_samples, n_classes))
        else:
            Y = xp.zeros((n_samples, n_classes), dtype=data.dtype, device=device_)

            for i in range(n_samples):
                start = int(indptr[i])
                end = int(indptr[i + 1])
                for j in range(start, end):
                    Y[i, int(indices[j])] = data[j]

    elif y_type == "multilabel-indicator":
        if sparse_output:
            Y = sp.csr_matrix(y)
            if pos_label != 1:

Review comment (suggestion):
We have an `_isin` util that we might be able to use here. Moreover `isin` was also added to array-api-extra but not released yet

Code after:
        int_dtype_ = y.dtype
    else:
        int_dtype_ = indexing_dtype(xp)

    if y_type == "binary":
        if n_classes == 1:
            if sparse_output:
                return sp.csr_matrix((n_samples, 1), dtype=int)
            else:
                Y = xp.zeros((n_samples, 1), dtype=int_dtype_)
                Y += neg_label
                return Y
        elif n_classes >= 3:
            y_type = "multiclass"

    sorted_class = xp.sort(classes)
    if y_type == "multilabel-indicator":
        y_n_classes = y.shape[1] if hasattr(y, "shape") else len(y[0])
        if n_classes != y_n_classes:
            raise ValueError(
                "classes {0} mismatch with the labels {1} found in the data".format(
                    classes, unique_labels(y)
                )
            )

    if y_type in ("binary", "multiclass"):
        y = column_or_1d(y)

        # pick out the known labels from y
        y_in_classes = _isin(y, classes, xp=xp)
        y_seen = y[y_in_classes]
        indices = xp.searchsorted(sorted_class, y_seen)
        # cast `y_in_classes` to integer dtype for `xp.cumulative_sum`
        y_in_classes = xp.astype(y_in_classes, int_dtype_)
        indptr = xp.concat(
            (
                xp.asarray([0], device=device_),
                xp.cumulative_sum(y_in_classes, axis=0),
            )
        )
        data = xp.full_like(indices, pos_label)

        # Use NumPy to construct the sparse matrix of one-hot labels
        Y = sp.csr_matrix(
            (
                _convert_to_numpy(data, xp=xp),
                _convert_to_numpy(indices, xp=xp),
                _convert_to_numpy(indptr, xp=xp),
            ),
            shape=(n_samples, n_classes),
        )