Some observations

1. The dense and sparse computers both yield the same result when max(X.shape[0], Y.shape[0]) < chunk_size, otherwise the sparse computation is still incorrect
2. (unrelated but could be substantial) It seems that we don't handle the case when parameter > X.shape[0], in that scenario even the values of the dense computer are wrong.

Testing script:

This is working ✅

import numpy as np
from numpy.testing import assert_array_equal
from scipy.sparse import csr_matrix
from sklearn.metrics._pairwise_distances_reduction import ArgKmin

X = np.array([[0, 1], [0, 0], [2, 3]], dtype=np.float64)
Y = np.array([[1, 1], [0, 0], [0, 1]], dtype=np.float64)

X_csr = csr_matrix(X)
Y_csr = csr_matrix(Y)

parameter = 3
dist, indices = ArgKmin.compute(
    X,
    Y,
    parameter,
    chunk_size=100,
    return_distance=True,
)
dist_csr, indices_csr = ArgKmin.compute(
    X_csr,
    Y_csr,
    parameter,
    chunk_size=100,
    return_distance=True,
)
assert_array_equal(dist, dist_csr)

This fails ❌ (notice that the only difference is the length of X and Y)

import numpy as np
from numpy.testing import assert_array_equal
from scipy.sparse import csr_matrix
from sklearn.metrics._pairwise_distances_reduction import ArgKmin

X = np.array([[0, 1], [0, 0], [2, 3]] * 40, dtype=np.float64)
Y = np.array([[1, 1], [0, 0], [0, 1]] * 40, dtype=np.float64)

X_csr = csr_matrix(X)
Y_csr = csr_matrix(Y)

parameter = 3
dist, indices = ArgKmin.compute(
    X,
    Y,
    parameter,
    chunk_size=100,
    return_distance=True,
)
dist_csr, indices_csr = ArgKmin.compute(
    X_csr,
    Y_csr,
    parameter,
    chunk_size=100,
    return_distance=True,
)
assert_array_equal(dist, dist_csr)

(unrelated but could be substantial) It seems that we don't handle the case when parameter > X.shape[0], in that scenario even the values of the dense computer are wrong.

Yes, that's right (I supposed you meant "when parameter > Y.shape[0]"?): it's not validated in ArgKmin . Most validations are done closer to user-facing interfaces, here in kneighbors:

As those interface are directly used by users, we do not revalidate all the parameters but this might change in the future.

As most of this work (except squeuclidean_row_norms) is not user-facing, should we append it to the changelog?

Hardware Scalability on 76574eb (using this script)

The current implementation of Sparse Sparse Euclidean NearestNeighbors does not benefit from PairwiseDistanceReduction introduced in #22134.

Following the design guidelines previously established, this new implementation:

• Computes the sqeuclidean_row_norm from a CSR format without the GIL
• Computes the middle term -2 * X @ Y.T in the idea introduced here. SparseSparseMiddleTermComputer realises a matmul operation in the blas gemm fashion without the GIL, between two CSR matrices.

This new design scales easily up to 8 cores on my laptop, unlike the current implementation that fails to do so.

Below are the results obtained from running the script linked in the title, with the following parameters:

• size of X and Y (100,000 x n_features)
• dtype: np.float64
• density: 0.05 (5%)

This branch

Main

However, I also observe that the absolute amount of time taken by this implementation greatly depends on the number of features: in high dimension (> 500), this branch currently performs worse than main, so I bit more investigation is required here.

n_features = 50 (main is blue)

n_features = 100 (main is blue)

n_features = 500 (main is blue)

Raw data

@Micky774 could you give me your thoughts about this PR? :)

Congrats @Vincent-Maladiere! 🎵 🎉 👏

That is significant performance improvements for many estimators! 🎉
I hope you have learnt a bunch of elements of scikit-learn and of Cython development. 😄