✔️ Linting Passed

All linting checks passed. Your pull request is in excellent shape! ☀️

_{Generated for commit: 78b9612. Link to the linter CI: here}

My main comment is why don't we convert to numpy at the top of the function and then use the same numpy code as before. We end up converting to numpy array to create the sparse matrix anyway? I would expect that manipulating 1d arrays on GPU is not going to have a huge speed improvement compared to using numpy, maybe that's a naive assumption.

It seems like this kind of stuff has been discussed before, for example #28626 (comment), #30439 (comment), and #30439 (comment). Unfortunately, I have to admit that I don't have a good overview of it and whether there was an agreed general direction ...

My intuition is telling me that it would be easier in metrics to convert first to numpy and then to do the computation in numpy. I am fine with it not being "real Array API support". As an end-user I want to do a cross_val_score with the heavy computation using the GPU to have speed improvement (I guess .fit and maybe .predict) and not have any error. If GPU 1d arrays are converted to numpy arrays to compute metrics and that is not the bottleneck I don't care that much.

My main comment is why don't we convert to numpy at the top of the function and then use the same numpy code as before. We end up converting to numpy array to create the sparse matrix anyway? I would expect that manipulating 1d arrays on GPU is not going to have a huge speed improvement compared to using numpy, maybe that's a naive assumption.

I think if the speed benefits gained from using the array api are not significant within the computation then we can probably just use numpy from start to end. However if there is benefit in using the array api in places where it can be used then I think it would be useful to use the array api and convert to numpy only where required.

@lesteve and @OmarManzoor:

I have checked how it would be if we converted all the input into numpy in the param validation in the top of the function and then return the confusion_matrix in whatever namespace the input came from in a separate branch (not sure if you can see the diff) These are the results:

execution_time for device cpu: 0.6608354229999804
execution_time for device cuda: 0.8810625369999343

And here, for comparison again the results from the current branch:

execution_time for device cpu: 2.0540889530000186
execution_time for device cuda: 0.6095863729999564

I ran the same test as here, but larger arrays (size=(int(1e6),)).

It seems that the numpy version is faster for numpy inputs and there is not a big difference for cuda, so I would tend to prefer the new branch version. Would you agree or is there something else to consider?

It seems that the numpy version is faster for numpy inputs and there is not a big difference for cuda, so I would tend to prefer the new branch version. Would you agree or is there something else to consider?

Looking at the results I think we might just convert to numpy and use the original code instead of complicating the logic inside this function with the array api.

I also ran some benchmarks on a kaggle notebook using arrays of size 1e8 and the results show a similar trend. Even though the CUDA version seems to be somewhat better consistently by using the array api, I don't think the performance gain is too dramatic.

RAM 29GB
GPU P100 16GB
array size 1e8

simple numpy by converting to numpy at the start

avg execution_time for device cpu: 7.236887979507446
avg execution_time for device cuda: 5.599829816818238

current code with array api

avg execution_time for device cpu: 11.78741683959961
avg execution_time for device cuda: 3.5200899124145506

Here is the short python script that I used

from time import time

import torch
from tqdm import tqdm

from sklearn.base import config_context
from sklearn.metrics import confusion_matrix

xp = torch

execution_times = []
for _ in tqdm(range(10), desc="CPU"):
    y_true = xp.randint(low=0, high=10, size=(int(1e8),), dtype=xp.int64, device="cpu")
    y_pred = xp.randint(low=0, high=10, size=(int(1e8),), dtype=xp.int64, device="cpu")
    start = time()
    with config_context(array_api_dispatch=True):
        confusion_matrix(y_true, y_pred)
    execution_times.append(time() - start)
print(f"avg execution_time for device cpu: {sum(execution_times) / 10}")

execution_times = []
for _ in tqdm(range(10), desc="CUDA"):
    y_true = xp.randint(low=0, high=10, size=(int(1e8),), dtype=xp.int64, device="cuda")
    y_pred = xp.randint(low=0, high=10, size=(int(1e8),), dtype=xp.int64, device="cuda")
    start = time()
    with config_context(array_api_dispatch=True):
        confusion_matrix(y_true, y_pred)
    execution_times.append(time() - start)
print(f"execution_time for device cuda: {sum(execution_times) / 10}")

It seems that the numpy version is faster for numpy inputs and there is not a big difference for cuda, so I would tend to prefer the new branch version. Would you agree or is there something else to consider?

I agree, but I would not have expected that the NumPy or torch CPU performance to be that degraded in this branch compared to the other branch. This is visible both in your int(1e6) data benchmark and in @OmarManzoor's larger data benchmark above (#30440 (comment)).

It might be interesting to profile both branches with a large enough dataset to understand where this discrepancy comes from.

Note that in the other branch, you move back the results to the input arrays' device:

https://github.com/StefanieSenger/scikit-learn/blob/1f23f634d55a9f9b7182584ddb6e6afd37548d98/sklearn/metrics/_classification.py#L424

I am not sure if this is needed / useful:

• on the plus side, it is less surprising to the caller;
• on the minus side, it might cause unnecessary data transfer / synchronization overhead between the host and the device: in general confusion matrices are tiny matrices (n_classes is seldom larger than 1000) so I don't expect downstream computation to ever benefit from GPU computation, even when the inputs y_true and y_pred are huge arrays.

Shall we even keep the result as a NumPy array?

I did some profiling using py-spy and here is what I found:

• when using a very large number of data points (int(1e8)), computing unique input values in y_true and y_pred dominates (more than 40% of the total time), even for numpy inputs that benefit from the unique value cache (sklearn.utils._unique.attach_unique/cached_unique);
• the numpy-dtype based unique cache system does not work when using pytorch tensors (as expected...) and then computing unique values takes 80% or more of the total time because it happens several times in _check_targets and then in unique_labels;
• the measured times vary a lot from run to run, and I am not sure if there is a significant difference between the two branches for numpy inputs.

However, converting early to numpy should allow benefiting from the caching of unique values if we convert to numpy before calling attach_unique.

@StefanieSenger could you please open a PR from your other branch (and cross-link to here) and then move the calls to attach_unique after the numpy conversion at the beginning of:

https://github.com/StefanieSenger/scikit-learn/blob/1f23f634d55a9f9b7182584ddb6e6afd37548d98/sklearn/metrics/_classification.py#L340-L367

• on the plus side, it is less surprising to the caller;
• on the minus side, it might cause unnecessary data transfer / synchronization overhead between the host and the device: in general confusion matrices are tiny matrices (n_classes is seldom larger than 1000) so I don't expect downstream computation to ever benefit from GPU computation, even when the inputs y_true and y_pred are huge arrays.

Shall we even keep the result as a NumPy array?

I think we can keep the result as a NumPy array as this returns a metric which I doubt would be utilized in further computations that require a GPU specifically.

Shall we even keep the result as a NumPy array?

I think we can keep the result as a NumPy array as this returns a metric which I doubt would be utilized in further computations that require a GPU specifically.

This is interesting and quite surprising to me, have you included this pattern anywhere else in scikit-learn? In SciPy, we have been quite careful to ensure that everything is namespace_in === namespace_out. I think there is a concern that telling users "mostly namespace_in === namespace_out, except ..." will be quite confusing. What if a user wants to do this?

# arrays of some namespace xp
x1 = confusion_matrix(y_true,y_pred)
z = xp.concat((x1, x2))

Since xp.concat may only accept xp arrays and not numpy arrays, you'd have to make sure users know to call xp.asarray on the confusion matrix. I think it would be pretty confusing if you have to do this for some sklearn functions and not others.

I think that adding an option to short-circuit and return the result as a NumPy array sounds like a great idea, but I don't think it should be the default behaviour.

I think there is a concern that telling users "mostly namespace_in === namespace_out, except ..." will be quite confusing.

I agree. I also think that converting back to the namespace that was passed, wouldn't be very costly here.

I think that adding an option to short-circuit and return the result as a NumPy array sounds like a great idea, but I don't think it should be the default behaviour.

So far, we did not have to change the public API of scikit-learn to add array API support (beyond enabling/disabling array API dispatch via sklearn.config_context or sklearn.set_config). Since confusion_matrix is a public function, adding such an option would set a precedent. I would therefore only add such an argument to a private _confusion_matrix function meant to be called either by the public confusion_matrix function or directly by public metric functions that know whether it's useful to keep the result as a numpy array.

Since confusion_matrix is a public function, adding such an option would set a precedent. I would therefore only add such an argument to a private _confusion_matrix function meant to be called either by the public confusion_matrix function or directly by public metric functions that know whether it's useful to keep the result as a numpy array.

In that case wouldn't it be just better to cast just before returning to the intended namespace and device?