scikit-learn · jjerphan · Sep 21, 2022 · Feb 20, 2022 · Feb 20, 2022 · Feb 22, 2022
diff --git a/doc/conftest.py b/doc/conftest.py
@@ -107,6 +107,14 @@ def skip_if_matplotlib_not_installed(fname):
         raise SkipTest(f"Skipping doctests for {basename}, matplotlib not installed")
 
 
+def skip_if_cupy_not_installed(fname):
+    try:
+        import cupy  # noqa
+    except ImportError:
+        basename = os.path.basename(fname)
+        raise SkipTest(f"Skipping doctests for {basename}, cupy not installed")
+
+
 def pytest_runtest_setup(item):
     fname = item.fspath.strpath
     # normalize filename to use forward slashes on Windows for easier handling
@@ -147,6 +155,9 @@ def pytest_runtest_setup(item):
         if fname.endswith(each):
             skip_if_matplotlib_not_installed(fname)
 
+    if fname.endswith("array_api.rst"):
+        skip_if_cupy_not_installed(fname)
+
 
 def pytest_configure(config):
     # Use matplotlib agg backend during the tests including doctests

diff --git a/doc/dispatching.rst b/doc/dispatching.rst
@@ -0,0 +1,14 @@
+.. Places parent toc into the sidebar
+
+:parenttoc: True
+
+.. include:: includes/big_toc_css.rst
+
+===========
+Dispatching
+===========
+
+.. toctree::
+    :maxdepth: 2
+
+    modules/array_api
diff --git a/doc/modules/array_api.rst b/doc/modules/array_api.rst
@@ -0,0 +1,75 @@
+.. Places parent toc into the sidebar
+
+:parenttoc: True
+
+.. _array_api:
+
+================================
+Array API support (experimental)
+================================
+
+.. currentmodule:: sklearn
+
+The `Array API <https://data-apis.org/array-api/latest/>`_ specification defines
+a standard API for all array manipulation libraries with a NumPy-like API.
+
+Some scikit-learn estimators that primarily rely on NumPy (as opposed to using
+Cython) to implement the algorithmic logic of their `fit`, `predict` or
+`transform` methods can be configured to accept any Array API compatible input
+datastructures and automatically dispatch operations to the underlying namespace
+instead of relying on NumPy.
+
+At this stage, this support is **considered experimental** and must be enabled
+explicitly as explained in the following.
+
+.. note::
+    Currently, only `cupy.array_api` and `numpy.array_api` are known to work
+    with scikit-learn's estimators.
+
+Example usage
+=============
+
+Here is an example code snippet to demonstrate how to use `CuPy
+<https://cupy.dev/>`_ to run
+:class:`~discriminant_analysis.LinearDiscriminantAnalysis` on a GPU::
+
+    >>> from sklearn.datasets import make_classification
+    >>> from sklearn import config_context
+    >>> from sklearn.discriminant_analysis import LinearDiscriminantAnalysis
+    >>> import cupy.array_api as xp
+
+    >>> X_np, y_np = make_classification(random_state=0)
+    >>> X_cu = xp.asarray(X_np)
+    >>> y_cu = xp.asarray(y_np)
+    >>> X_cu.device
+    <CUDA Device 0>
+
+    >>> with config_context(array_api_dispatch=True):
+    ...     lda = LinearDiscriminantAnalysis()
+    ...     X_trans = lda.fit_transform(X_cu, y_cu)
+    >>> X_trans.device
+    <CUDA Device 0>
+
+After the model is trained, fitted attributes that are arrays will also be
+from the same Array API namespace as the training data. For example, if CuPy's
+Array API namespace was used for training, then fitted attributes will be on the
+GPU. We provide a experimental `_estimator_with_converted_arrays` utility that
+transfers an estimator attributes from Array API to a ndarray::
+
+    >>> from sklearn.utils._array_api import _estimator_with_converted_arrays
+    >>> cupy_to_ndarray = lambda array : array._array.get()
+    >>> lda_np = _estimator_with_converted_arrays(lda, cupy_to_ndarray)
+    >>> X_trans = lda_np.transform(X_np)
+    >>> type(X_trans)
+    <class 'numpy.ndarray'>
+
+.. _array_api_estimators:
+
+Estimators with support for `Array API`-compatible inputs
+=========================================================
+
+- :class:`discriminant_analysis.LinearDiscriminantAnalysis` (with `solver="svd"`)
+
+Coverage for more estimators is expected to grow over time. Please follow the
+dedicated `meta-issue on GitHub
+<https://github.com/scikit-learn/scikit-learn/issues/22352>`_ to track progress.
diff --git a/doc/user_guide.rst b/doc/user_guide.rst
@@ -30,3 +30,4 @@ User Guide
    computing.rst
    model_persistence.rst
    common_pitfalls.rst
+   dispatching.rst
diff --git a/doc/whats_new/v1.2.rst b/doc/whats_new/v1.2.rst
@@ -191,6 +191,12 @@ Changelog
 :mod:`sklearn.discriminant_analysis`
 ....................................
 
+- |MajorFeature| :class:`discriminant_analysis.LinearDiscriminantAnalysis` now
+  supports the `Array API <https://data-apis.org/array-api/latest/>`_ for
+  `solver="svd"`. Array API support is considered experimental and might evolve
+  without being subjected to our usual rolling deprecation cycle policy. See
+  :ref:`array_api` for more details. :pr:`22554` by `Thomas Fan`_.
+
 - |Fix| Validate parameters only in `fit` and not in `__init__`
   for :class:`discriminant_analysis.QuadraticDiscriminantAnalysis`.
   :pr:`24218` by :user:`Stefanie Molin <stefmolin>`.

diff --git a/sklearn/_config.py b/sklearn/_config.py
@@ -13,6 +13,7 @@
         os.environ.get("SKLEARN_PAIRWISE_DIST_CHUNK_SIZE", 256)
     ),
     "enable_cython_pairwise_dist": True,
+    "array_api_dispatch": False,
 }
 _threadlocal = threading.local()
 
@@ -50,6 +51,7 @@ def set_config(
     display=None,
     pairwise_dist_chunk_size=None,
     enable_cython_pairwise_dist=None,
+    array_api_dispatch=None,
 ):
     """Set global scikit-learn configuration
 
@@ -110,6 +112,14 @@ def set_config(
 
         .. versionadded:: 1.1
 
+    array_api_dispatch : bool, default=None
+        Use Array API dispatching when inputs follow the Array API standard.
+        Default is False.
+
+        See the :ref:`User Guide <array_api>` for more details.
+
+        .. versionadded:: 1.2
+
     See Also
     --------
     config_context : Context manager for global scikit-learn configuration.
@@ -129,6 +139,8 @@ def set_config(
         local_config["pairwise_dist_chunk_size"] = pairwise_dist_chunk_size
     if enable_cython_pairwise_dist is not None:
         local_config["enable_cython_pairwise_dist"] = enable_cython_pairwise_dist
+    if array_api_dispatch is not None:
+        local_config["array_api_dispatch"] = array_api_dispatch
 
 
 @contextmanager
@@ -140,6 +152,7 @@ def config_context(
     display=None,
     pairwise_dist_chunk_size=None,
     enable_cython_pairwise_dist=None,
+    array_api_dispatch=None,
 ):
     """Context manager for global scikit-learn configuration.
 
@@ -199,6 +212,14 @@ def config_context(
 
         .. versionadded:: 1.1
 
+    array_api_dispatch : bool, default=None
+        Use Array API dispatching when inputs follow the Array API standard.
+        Default is False.
+
+        See the :ref:`User Guide <array_api>` for more details.
+
+        .. versionadded:: 1.2
+
     Yields
     ------
     None.
@@ -234,6 +255,7 @@ def config_context(
         display=display,
         pairwise_dist_chunk_size=pairwise_dist_chunk_size,
         enable_cython_pairwise_dist=enable_cython_pairwise_dist,
+        array_api_dispatch=array_api_dispatch,
     )
 
     try: