Path for pluggable low-level computational routines

The goal of this issue is to discuss the design and prototype a way to register alternative implementations for core low level routines in scikit-learn, in particular to benefit from hardware optimized implementations (e.g. using GPUs efficiently).

## Motivation

scikit-learn aims to provide reasonably easy to maintain and portable implementations of standard machine learning algorithms. Those implementations are typically written in Python (with the help of NumPy and SciPy) or in Cython when the overhead of the Python interpreter prevents us to efficiently implement algorithms with (nested) tight loops. This allows us reasonably fast implementations as binary packages (installable with pip/PyPI, conda/conda-forge, conda/anaconda or Linux distros) for a variety of platforms (Linux / macOS / Windows) x (x86_64, i686, arm64, ppc64le) from a single code-base with no external runtime dependencies beyond Python, NumPy and SciPy.

Recently, GPU hardware have proven very competitive for many machine learning related workloads, either from a pure latency standpoint, or from the standpoint of a better computation/energy trade-off (irrespective of the raw speed considerations). However, hardware optimized implementation are typically not portable and mandates additional dependencies.

We therefore propose to design a way for our users to register alternative implementations of low-level computation routines in scikit-learn provided they has installed the required extension package(s) that matches their specific hardware.

## Relationship to adopting the Array API spec

This proposal is related and complementary to another effort, namely:

- #22352

The Array API spec is useful to make it possible to have some scikit-learn estimators written using a pure numpy syntax to delegate the computation to alternative Array API compatible libraries such as CuPy.

However, some algorithms in scikit-learn cannot be efficiently written using NumPy operation only, for instance the main [K-Means loop is written in Cython](https://github.com/scikit-learn/scikit-learn/blob/0de3b0d1eaacee9f7b15cabc05752cba945c7621/sklearn/cluster/_k_means_lloyd.pyx#L116-L152) to process chunks of samples in parallel (using `prange` and OpenMP), compute distance with centroids and reduce those distances to find assign each sample to its closest centroid on-the-fly while preventing unnecessary memory transfer between CPU cache and RAM.

If we want to run this algorithm efficiently on GPU hardware, one would need to dispatch the computation of this low level function to an alternative implementation that can work on GPU, either written in C/C++ with GPU-specific supporting runtime libraries and compilers (e.g. OpenCL, NVIDIA Cuda, Intel oneAPI DPC++, AMD ROCm...)  or using a Python syntax with the help of GPU support provided in numba for instance.

## List of candidate routines

- the main k-means loop (pairwise distances between samples and centroids with on-the-fly argmin reduction) 
- the core k-nearest neighbors computation loop (pairwise distances with on-the-fly arg-k-min reduction) 
- pairwise distances (without arg-k-min reduction)
- pairwise kernels computation (e.g. for use in Nystroem)
- ...

## Explicit registration API design ideas

I started to draft some API ideas in:

- https://hackmd.io/@4rHCRgfySZSdd5eMtfUJiA/S1RDJ3HCF

Feel free to comment here, or there.

This design is subject to evolve, in particular to make it possible to register both Array API extensions and non-Array API extensions with the same registration API.

## Next steps

- start a proof of concept implementation in a branch + external repo using either numba or C++ for k-means. Edit: this is happening here:
  - #25535
  - https://github.com/soda-inria/sklearn-numba-dpex
- once the API design has converged on the PoC, write a formal SLEP.

COMMENT:
Here are is a proposal on the developer API for the multi-dispatching:

## scikit-learn dispatching

This registers `pairwise_distance` as a function that can be dispatched:

```python
from sklearn.utils.dispatch import multidispatch

# Dispatch on "X" and "metrics"
@multidispatch(domain="sklearn.metrics", keys=["X", "metric"])
def pairwise_distance(...):
    ...
```

## Backend Provider

### Using Python types

The following registers `faster_pairwise` and accepts the types specified by the Python types. This backend only supports "euclidean" and "manhattan" metrics:

```python
from sklearn.utils.dispatch import Backend
import numpy as np

be = Backend(domain="sklearn.metrics")

@be.register(name="pairwise_distance")
def faster_pairwise(X: np.ndarray, .., metric: Literal["euclidean", "manhattan"]):
    ...
```

### Specifying a CuPy array:

Specifying `cp.ndarray` in the typing annotation dispatches for CuPy arrays:

```python
import cupy as cp

@be.register(name="pairwise_distance")
def faster_pairwise(X: cp.ndarray, ...):
    ...
```

### Backend only supports float32

There can be a callable that is passed in the dispatch keys and returns True if the function supports the input. 

```python
# Uses a **kwargs just in case we want to dispatch more keys in the future
@be.register(
    name="pairwise_distance",
    supports=lambda **kwargs: kwargs["X"].dtype == np.float32
)
def faster_pairwise(X: np.ndarray, ...):
    ...
```

Note that this callable can be used instead of Python typing.

Labels: API, Performance