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

Code before:
        y = column_or_1d(y)

        # pick out the known labels from y
        y_in_classes = _isin(y, classes, 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)

        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:
                data = xp.full_like(Y.data, pos_label)
                Y.data = data

        else:
            if sp.issparse(y):
                y = y.toarray()

            Y = xp.asarray(y, device=device_, copy=True)
            if pos_label != 1:
                Y[Y != 0] = pos_label

    else:
        raise ValueError(
            "%s target data is not supported with label binarization" % y_type
        )

    if not sparse_output:
        if neg_label != 0:
            Y[Y == 0] = neg_label


Review comment (suggestion):
I did some benchmarking and I don't see improvements when using CUDA:

```bash
Average `label_binarize` Runtime Result for Various Array API:
Number of Trials: 100    Numbere of Samples: 100000      Number of Classes: 5000
NumPy: 0.878273 seconds
PyTorch: 2.454254 seconds, speed-up: 0.357858x
```
<details>

<summary>

*Code to Reproduce*

</summary>

```pyrhon
from time import time

import numpy as np
import torch
from tqdm import tqdm

from sklearn import config_context
from sklearn.preprocessing._label import label_binarize

N_CLASSES = 5000
N_SAMPLES = 100000
N_TRIALS = 100

y_np = np.random.randint(0, N_CLASSES, size=N_SAMPLES)
classes = [i for i in range(N_CLASSES)]

# NumPy Benchmarks
fit_times = []
for _ in tqdm(range(N_TRIALS), desc="Benchmarking for NumPy"):
    start = time()
    label_binarize(y=y_np, classes=classes)
    fit_times.append(time() - start)

numpy_average = sum(fit_times) / N_TRIALS

# Torch CUDA Benchmarks
y_tensor = torch.tensor(y_np, device="cuda")
fit_times = []
with config_context(array_api_dispatch=True):
    for _ in tqdm(range(N_TRIALS), desc="Benchmarking for Torch CUDA"):
        start = time()
        label_binarize(y=y_tensor, classes=classes)
        torch.cuda.synchronize()
        fit_times.append(time() - start)

torch_average = sum(fit_times) / N_TRIALS

print("Average `label_binarize` Runtime Result for Various Array API:")
print(
    f"Number of Trials: {N_TRIALS} \t "
    f"Numbere of Samples: {N_SAMPLES} \t Number of Classes: {N_CLASSES}"
)
print(f"NumPy: {numpy_average:.6f} seconds")
print(
    f"PyTorch: {torch_average:.6f} seconds, "
    f"speed-up: {numpy_average / torch_average:.6f}x"
)
```
</details>

I suspect the issue comes from the part where we construct the dense binary matrix using this double `for`-loop. This is likely because `array-api-strict` doesn't support "fancy indexing":

```python
Y[i, indices[indptr[i]:indptr[i+1]]] = data[indptr[i]:indptr[i+1]]
```

I'll investigate this further.

Code after:
    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),
        )

        if not sparse_output:
            Y = xp.asarray(Y.toarray(), device=device_)

    elif y_type == "multilabel-indicator":
        if sparse_output:
            Y = sp.csr_matrix(y)
            if pos_label != 1:
                data = xp.full_like(Y.data, pos_label)
                Y.data = data
        else:
            if sp.issparse(y):
                y = y.toarray()

            Y = xp.asarray(y, device=device_, copy=True)
            if pos_label != 1:
                Y[Y != 0] = pos_label

    else:
        raise ValueError(
            "%s target data is not supported with label binarization" % y_type
        )

    if not sparse_output:
        if neg_label != 0: