scikit-learn · betatim · Apr 21, 2023 · Mar 23, 2023 · Mar 23, 2023 · Mar 23, 2023
diff --git a/build_tools/azure/pylatest_conda_forge_mkl_linux-64_conda.lock b/build_tools/azure/pylatest_conda_forge_mkl_linux-64_conda.lock
diff --git a/build_tools/azure/pylatest_conda_forge_mkl_linux-64_environment.yml b/build_tools/azure/pylatest_conda_forge_mkl_linux-64_environment.yml
@@ -20,3 +20,6 @@ dependencies:
   - pytest-cov
   - coverage
   - ccache
+  - pytorch=1.13
+  - pytorch-cpu
+  - array-api-compat
diff --git a/build_tools/update_environments_and_lock_files.py b/build_tools/update_environments_and_lock_files.py
@@ -88,9 +88,15 @@ def remove_from(alist, to_remove):
         "folder": "build_tools/azure",
         "platform": "linux-64",
         "channel": "conda-forge",
-        "conda_dependencies": common_dependencies + ["ccache"],
+        "conda_dependencies": common_dependencies + [
+            "ccache",
+            "pytorch",
+            "pytorch-cpu",
+            "array-api-compat",
+        ],
         "package_constraints": {
             "blas": "[build=mkl]",
+            "pytorch": "1.13",
         },
     },
     {

diff --git a/doc/modules/array_api.rst b/doc/modules/array_api.rst
@@ -12,6 +12,8 @@ Array API support (experimental)
 
 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.
+Scikit-learn's Array API support requires
+`array-api-compat <https://github.com/data-apis/array-api-compat>`__ to be installed.
 
 Some scikit-learn estimators that primarily rely on NumPy (as opposed to using
 Cython) to implement the algorithmic logic of their `fit`, `predict` or
@@ -23,8 +25,8 @@ 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.
+    Currently, only `cupy.array_api`, `numpy.array_api`, `cupy`, and `PyTorch`
+    are known to work with scikit-learn's estimators.
 
 Example usage
 =============
@@ -36,11 +38,11 @@ Here is an example code snippet to demonstrate how to use `CuPy
     >>> from sklearn.datasets import make_classification
     >>> from sklearn import config_context
     >>> from sklearn.discriminant_analysis import LinearDiscriminantAnalysis
-    >>> import cupy.array_api as xp
+    >>> import cupy
 
     >>> X_np, y_np = make_classification(random_state=0)
-    >>> X_cu = xp.asarray(X_np)
-    >>> y_cu = xp.asarray(y_np)
+    >>> X_cu = cupy.asarray(X_np)
+    >>> y_cu = cupy.asarray(y_np)
     >>> X_cu.device
     <CUDA Device 0>
 
@@ -57,12 +59,30 @@ 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()
+    >>> cupy_to_ndarray = lambda 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'>
 
+PyTorch Support
+---------------
+
+PyTorch Tensors are supported by setting `array_api_dispatch=True` and passing in
+the tensors directly::
+
+    >>> import torch
+    >>> X_torch = torch.asarray(X_np, device="cuda", dtype=torch.float32)
+    >>> y_torch = torch.asarray(y_np, device="cuda", dtype=torch.float32)
+
+    >>> with config_context(array_api_dispatch=True):
+    ...     lda = LinearDiscriminantAnalysis()
+    ...     X_trans = lda.fit_transform(X_torch, y_torch)
+    >>> type(X_trans)
+    <class 'torch.Tensor'>
+    >>> X_trans.device.type
+    'cuda'
+
 .. _array_api_estimators:
 
 Estimators with support for `Array API`-compatible inputs

diff --git a/doc/whats_new/v1.3.rst b/doc/whats_new/v1.3.rst
@@ -230,6 +230,13 @@ Changelog
   :class:`decomposition.MiniBatchNMF` which can produce different results than previous
   versions. :pr:`25438` by :user:`Yotam Avidar-Constantini <yotamcons>`.
 
+:mod:`sklearn.discriminant_analysis`
+....................................
+
+- |Enhancement| :class:`discriminant_analysis.LinearDiscriminantAnalysis` now
+  supports the `PyTorch <https://pytorch.org/>`__. See
+  :ref:`array_api` for more details. :pr:`25956` by `Thomas Fan`_.
+
 :mod:`sklearn.ensemble`
 .......................
 

diff --git a/sklearn/_config.py b/sklearn/_config.py
@@ -154,6 +154,9 @@ def set_config(
     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:
+        from .utils._array_api import _check_array_api_dispatch
+
+        _check_array_api_dispatch(array_api_dispatch)
         local_config["array_api_dispatch"] = array_api_dispatch
     if transform_output is not None:
         local_config["transform_output"] = transform_output

diff --git a/sklearn/discriminant_analysis.py b/sklearn/discriminant_analysis.py
@@ -21,7 +21,7 @@
 from .covariance import ledoit_wolf, empirical_covariance, shrunk_covariance
 from .utils.multiclass import unique_labels
 from .utils.validation import check_is_fitted
-from .utils._array_api import get_namespace, _expit
+from .utils._array_api import get_namespace, _expit, device, size
 from .utils.multiclass import check_classification_targets
 from .utils.extmath import softmax
 from .utils._param_validation import StrOptions, Interval, HasMethods
@@ -107,11 +107,11 @@ def _class_means(X, y):
     means : array-like of shape (n_classes, n_features)
         Class means.
     """
-    xp, is_array_api = get_namespace(X)
+    xp, is_array_api_compliant = get_namespace(X)
     classes, y = xp.unique_inverse(y)
-    means = xp.zeros(shape=(classes.shape[0], X.shape[1]))
+    means = xp.zeros((classes.shape[0], X.shape[1]), device=device(X), dtype=X.dtype)
 
-    if is_array_api:
+    if is_array_api_compliant:
         for i in range(classes.shape[0]):
             means[i, :] = xp.mean(X[y == i], axis=0)
     else:
@@ -483,9 +483,9 @@ def _solve_svd(self, X, y):
         y : array-like of shape (n_samples,) or (n_samples, n_targets)
             Target values.
         """
-        xp, is_array_api = get_namespace(X)
+        xp, is_array_api_compliant = get_namespace(X)
 
-        if is_array_api:
+        if is_array_api_compliant:
             svd = xp.linalg.svd
         else:
             svd = scipy.linalg.svd
@@ -586,9 +586,9 @@ def fit(self, X, y):
 
         if self.priors is None:  # estimate priors from sample
             _, cnts = xp.unique_counts(y)  # non-negative ints
-            self.priors_ = xp.astype(cnts, xp.float64) / float(y.shape[0])
+            self.priors_ = xp.astype(cnts, X.dtype) / float(y.shape[0])
         else:
-            self.priors_ = xp.asarray(self.priors)
+            self.priors_ = xp.asarray(self.priors, dtype=X.dtype)
 
         if xp.any(self.priors_ < 0):
             raise ValueError("priors must be non-negative")
@@ -634,7 +634,7 @@ def fit(self, X, y):
                 shrinkage=self.shrinkage,
                 covariance_estimator=self.covariance_estimator,
             )
-        if self.classes_.size == 2:  # treat binary case as a special case
+        if size(self.classes_) == 2:  # treat binary case as a special case
             coef_ = xp.asarray(self.coef_[1, :] - self.coef_[0, :], dtype=X.dtype)
             self.coef_ = xp.reshape(coef_, (1, -1))
             intercept_ = xp.asarray(
@@ -688,9 +688,9 @@ def predict_proba(self, X):
             Estimated probabilities.
         """
         check_is_fitted(self)
-        xp, is_array_api = get_namespace(X)
+        xp, is_array_api_compliant = get_namespace(X)
         decision = self.decision_function(X)
-        if self.classes_.size == 2:
+        if size(self.classes_) == 2:
             proba = _expit(decision)
             return xp.stack([1 - proba, proba], axis=1)
         else:

diff --git a/sklearn/linear_model/_base.py b/sklearn/linear_model/_base.py
@@ -449,7 +449,7 @@ def predict(self, X):
         else:
             indices = xp.argmax(scores, axis=1)
 
-        return xp.take(self.classes_, indices, axis=0)
+        return xp.take(self.classes_, indices)
 
     def _predict_proba_lr(self, X):
         """Probability estimation for OvR logistic regression.

diff --git a/sklearn/tests/test_config.py b/sklearn/tests/test_config.py
@@ -1,9 +1,11 @@
+import builtins
 import time
 from concurrent.futures import ThreadPoolExecutor
 
 import pytest
 
 from sklearn import get_config, set_config, config_context
+import sklearn
 from sklearn.utils.parallel import delayed, Parallel
 
 
@@ -145,3 +147,45 @@ def test_config_threadsafe():
         ]
 
     assert items == [False, True, False, True]
+
+
+def test_config_array_api_dispatch_error(monkeypatch):
+    """Check error is raised when array_api_compat is not installed."""
+
+    # Hide array_api_compat import
+    orig_import = builtins.__import__
+
+    def mocked_import(name, *args, **kwargs):
+        if name == "array_api_compat":
+            raise ImportError
+        return orig_import(name, *args, **kwargs)
+
+    monkeypatch.setattr(builtins, "__import__", mocked_import)
+
+    with pytest.raises(ImportError, match="array_api_compat is required"):
+        with config_context(array_api_dispatch=True):
+            pass
+
+    with pytest.raises(ImportError, match="array_api_compat is required"):
+        set_config(array_api_dispatch=True)
+
+
+def test_config_array_api_dispatch_error_numpy(monkeypatch):
+    """Check error when NumPy is too old"""
+    # Pretend that array_api_compat is installed.
+    orig_import = builtins.__import__
+
+    def mocked_import(name, *args, **kwargs):
+        if name == "array_api_compat":
+            return object()
+        return orig_import(name, *args, **kwargs)
+
+    monkeypatch.setattr(builtins, "__import__", mocked_import)
+    monkeypatch.setattr(sklearn.utils._array_api.numpy, "__version__", "1.20")
+
+    with pytest.raises(ImportError, match="NumPy must be 1.21 or newer"):
+        with config_context(array_api_dispatch=True):
+            pass
+
+    with pytest.raises(ImportError, match="NumPy must be 1.21 or newer"):
+        set_config(array_api_dispatch=True)
diff --git a/sklearn/tests/test_discriminant_analysis.py b/sklearn/tests/test_discriminant_analysis.py
@@ -13,6 +13,7 @@
 from sklearn.utils._testing import assert_almost_equal
 from sklearn.utils._array_api import _convert_to_numpy
 from sklearn.utils._testing import _convert_container
+from sklearn.utils._testing import skip_if_array_api_compat_not_configured
 
 from sklearn.datasets import make_blobs
 from sklearn.discriminant_analysis import LinearDiscriminantAnalysis
@@ -676,6 +677,7 @@ def test_get_feature_names_out():
     assert_array_equal(names_out, expected_names_out)
 
 
+@skip_if_array_api_compat_not_configured
 @pytest.mark.parametrize("array_namespace", ["numpy.array_api", "cupy.array_api"])
 def test_lda_array_api(array_namespace):
     """Check that the array_api Array gives the same results as ndarrays."""
@@ -725,6 +727,66 @@ def test_lda_array_api(array_namespace):
 
         result_xp_np = _convert_to_numpy(result_xp, xp=xp)
 
+        assert_allclose(
+            result,
+            result_xp_np,
+            err_msg=f"{method} did not the return the same result",
+            atol=1e-5,
+        )
+
+
+@skip_if_array_api_compat_not_configured
+@pytest.mark.parametrize("device", ["cuda", "cpu"])
+@pytest.mark.parametrize("dtype", ["float32", "float64"])
+def test_lda_array_torch(device, dtype):
+    """Check running on PyTorch Tensors gives the same results as NumPy"""
+    torch = pytest.importorskip("torch")
+    if device == "cuda" and not torch.has_cuda:
+        pytest.skip("test requires cuda")
+
+    lda = LinearDiscriminantAnalysis()
+    X_np = X6.astype(dtype)
+    y_np = y6.astype(dtype)
+    lda.fit(X_np, y_np)
+
+    X_torch = torch.asarray(X_np, device=device)
+    y_torch = torch.asarray(y_np, device=device)
+    lda_xp = clone(lda)
+    with config_context(array_api_dispatch=True):
+        lda_xp.fit(X_torch, y_torch)
+
+    array_attributes = {
+        key: value for key, value in vars(lda).items() if isinstance(value, np.ndarray)
+    }
+
+    for key, attribute in array_attributes.items():
+        lda_xp_param = getattr(lda_xp, key)
+        assert isinstance(lda_xp_param, torch.Tensor)
+        assert lda_xp_param.device.type == device
+
+        lda_xp_param_np = _convert_to_numpy(lda_xp_param, xp=torch)
+        assert_allclose(
+            attribute, lda_xp_param_np, err_msg=f"{key} not the same", atol=1e-3
+        )
+
+    # Check predictions are the same
+    methods = (
+        "decision_function",
+        "predict",
+        "predict_log_proba",
+        "predict_proba",
+        "transform",
+    )
+    for method in methods:
+        result = getattr(lda, method)(X_np)
+        with config_context(array_api_dispatch=True):
+            result_xp = getattr(lda_xp, method)(X_torch)
+
+        assert isinstance(result_xp, torch.Tensor)
+        assert result_xp.device.type == device
+
+        result_xp_np = _convert_to_numpy(result_xp, xp=torch)
+
         assert_allclose(
             result,
             result_xp_np,