scikit-learn · thomasjpfan · Feb 21, 2020 · Feb 21, 2020 · Feb 21, 2020 · Feb 21, 2020
diff --git a/azure-pipelines.yml b/azure-pipelines.yml
@@ -67,23 +67,23 @@ jobs:
         SKLEARN_SKIP_NETWORK_TESTS: '0'
 
 # Will run all the time regardless of linting outcome.
-- template: build_tools/azure/posix.yml
-  parameters:
-    name: Linux_Runs
-    vmImage: ubuntu-18.04
-    matrix:
-      pylatest_conda_mkl:
-        DISTRIB: 'conda'
-        PYTHON_VERSION: '*'
-        BLAS: 'mkl'
-        NUMPY_VERSION: '*'
-        SCIPY_VERSION: '*'
-        CYTHON_VERSION: '*'
-        PILLOW_VERSION: '*'
-        PYTEST_VERSION: '*'
-        JOBLIB_VERSION: '*'
-        THREADPOOLCTL_VERSION: '2.0.0'
-        COVERAGE: 'true'
+# - template: build_tools/azure/posix.yml
+#   parameters:
+#     name: Linux_Runs
+#     vmImage: ubuntu-18.04
+#     matrix:
+#       pylatest_conda_mkl:
+#         DISTRIB: 'conda'
+#         PYTHON_VERSION: '*'
+#         BLAS: 'mkl'
+#         NUMPY_VERSION: '*'
+#         SCIPY_VERSION: '*'
+#         CYTHON_VERSION: '*'
+#         PILLOW_VERSION: '*'
+#         PYTEST_VERSION: '*'
+#         JOBLIB_VERSION: '*'
+#         THREADPOOLCTL_VERSION: '2.0.0'
+#         COVERAGE: 'true'
 
 - template: build_tools/azure/posix.yml
   parameters:
@@ -95,31 +95,31 @@ jobs:
       # Linux environment to test that scikit-learn can be built against
       # versions of numpy, scipy with ATLAS that comes with Ubuntu Bionic 18.04
       # i.e. numpy 1.13.3 and scipy 0.19
-      py36_ubuntu_atlas:
-        DISTRIB: 'ubuntu'
-        PYTHON_VERSION: '3.6'
-        JOBLIB_VERSION: '0.11'
-        PYTEST_XDIST: 'false'
-        THREADPOOLCTL_VERSION: '2.0.0'
-      # Linux + Python 3.6 build with OpenBLAS and without SITE_JOBLIB
-      py36_conda_openblas:
-        DISTRIB: 'conda'
-        PYTHON_VERSION: '3.6'
-        BLAS: 'openblas'
-        NUMPY_VERSION: '1.13.3'
-        SCIPY_VERSION: '0.19.1'
-        PANDAS_VERSION: '*'
-        CYTHON_VERSION: '*'
-        # temporary pin pytest due to unknown failure with pytest 5.3
-        PYTEST_VERSION: '5.2'
-        PILLOW_VERSION: '4.2.1'
-        MATPLOTLIB_VERSION: '2.1.1'
-        SCIKIT_IMAGE_VERSION: '*'
-        # latest version of joblib available in conda for Python 3.6
-        JOBLIB_VERSION: '0.13.2'
-        THREADPOOLCTL_VERSION: '2.0.0'
-        PYTEST_XDIST: 'false'
-        COVERAGE: 'true'
+      # py36_ubuntu_atlas:
+      #   DISTRIB: 'ubuntu'
+      #   PYTHON_VERSION: '3.6'
+      #   JOBLIB_VERSION: '0.11'
+      #   PYTEST_XDIST: 'false'
+      #   THREADPOOLCTL_VERSION: '2.0.0'
+      # # Linux + Python 3.6 build with OpenBLAS and without SITE_JOBLIB
+      # py36_conda_openblas:
+      #   DISTRIB: 'conda'
+      #   PYTHON_VERSION: '3.6'
+      #   BLAS: 'openblas'
+      #   NUMPY_VERSION: '1.13.3'
+      #   SCIPY_VERSION: '0.19.1'
+      #   PANDAS_VERSION: '*'
+      #   CYTHON_VERSION: '*'
+      #   # temporary pin pytest due to unknown failure with pytest 5.3
+      #   PYTEST_VERSION: '5.2'
+      #   PILLOW_VERSION: '4.2.1'
+      #   MATPLOTLIB_VERSION: '2.1.1'
+      #   SCIKIT_IMAGE_VERSION: '*'
+      #   # latest version of joblib available in conda for Python 3.6
+      #   JOBLIB_VERSION: '0.13.2'
+      #   THREADPOOLCTL_VERSION: '2.0.0'
+      #   PYTEST_XDIST: 'false'
+      #   COVERAGE: 'true'
       # Linux environment to test the latest available dependencies and MKL.
       # It runs tests requiring lightgbm, pandas and PyAMG.
       pylatest_pip_openblas_pandas:
@@ -131,66 +131,66 @@ jobs:
         TEST_DOCSTRINGS: 'true'
         CHECK_WARNINGS: 'true'
 
-- template: build_tools/azure/posix-32.yml
-  parameters:
-    name: Linux32
-    vmImage: ubuntu-18.04
-    dependsOn: [linting]
-    condition: and(ne(variables['Build.Reason'], 'Schedule'), succeeded('linting'))
-    matrix:
-      py36_ubuntu_atlas_32bit:
-        DISTRIB: 'ubuntu-32'
-        PYTHON_VERSION: '3.6'
-        JOBLIB_VERSION: '0.13'
-        THREADPOOLCTL_VERSION: '2.0.0'
+# - template: build_tools/azure/posix-32.yml
+#   parameters:
+#     name: Linux32
+#     vmImage: ubuntu-18.04
+#     dependsOn: [linting]
+#     condition: and(ne(variables['Build.Reason'], 'Schedule'), succeeded('linting'))
+#     matrix:
+#       py36_ubuntu_atlas_32bit:
+#         DISTRIB: 'ubuntu-32'
+#         PYTHON_VERSION: '3.6'
+#         JOBLIB_VERSION: '0.13'
+#         THREADPOOLCTL_VERSION: '2.0.0'
 
-- template: build_tools/azure/posix.yml
-  parameters:
-    name: macOS
-    vmImage: macOS-10.14
-    dependsOn: [linting]
-    condition: and(ne(variables['Build.Reason'], 'Schedule'), succeeded('linting'))
-    matrix:
-      pylatest_conda_mkl:
-        DISTRIB: 'conda'
-        PYTHON_VERSION: '*'
-        BLAS: 'mkl'
-        NUMPY_VERSION: '*'
-        SCIPY_VERSION: '*'
-        CYTHON_VERSION: '*'
-        PILLOW_VERSION: '*'
-        PYTEST_VERSION: '*'
-        JOBLIB_VERSION: '*'
-        THREADPOOLCTL_VERSION: '2.0.0'
-        COVERAGE: 'true'
-      pylatest_conda_mkl_no_openmp:
-        DISTRIB: 'conda'
-        PYTHON_VERSION: '*'
-        BLAS: 'mkl'
-        NUMPY_VERSION: '*'
-        SCIPY_VERSION: '*'
-        CYTHON_VERSION: '*'
-        PILLOW_VERSION: '*'
-        PYTEST_VERSION: '*'
-        JOBLIB_VERSION: '*'
-        THREADPOOLCTL_VERSION: '2.0.0'
-        COVERAGE: 'true'
-        SKLEARN_TEST_NO_OPENMP: 'true'
-        SKLEARN_SKIP_OPENMP_TEST: 'true'
+# - template: build_tools/azure/posix.yml
+#   parameters:
+#     name: macOS
+#     vmImage: macOS-10.14
+#     dependsOn: [linting]
+#     condition: and(ne(variables['Build.Reason'], 'Schedule'), succeeded('linting'))
+#     matrix:
+#       pylatest_conda_mkl:
+#         DISTRIB: 'conda'
+#         PYTHON_VERSION: '*'
+#         BLAS: 'mkl'
+#         NUMPY_VERSION: '*'
+#         SCIPY_VERSION: '*'
+#         CYTHON_VERSION: '*'
+#         PILLOW_VERSION: '*'
+#         PYTEST_VERSION: '*'
+#         JOBLIB_VERSION: '*'
+#         THREADPOOLCTL_VERSION: '2.0.0'
+#         COVERAGE: 'true'
+#       pylatest_conda_mkl_no_openmp:
+#         DISTRIB: 'conda'
+#         PYTHON_VERSION: '*'
+#         BLAS: 'mkl'
+#         NUMPY_VERSION: '*'
+#         SCIPY_VERSION: '*'
+#         CYTHON_VERSION: '*'
+#         PILLOW_VERSION: '*'
+#         PYTEST_VERSION: '*'
+#         JOBLIB_VERSION: '*'
+#         THREADPOOLCTL_VERSION: '2.0.0'
+#         COVERAGE: 'true'
+#         SKLEARN_TEST_NO_OPENMP: 'true'
+#         SKLEARN_SKIP_OPENMP_TEST: 'true'
 
-- template: build_tools/azure/windows.yml
-  parameters:
-    name: Windows
-    vmImage: vs2017-win2016
-    dependsOn: [linting]
-    condition: and(ne(variables['Build.Reason'], 'Schedule'), succeeded('linting'))
-    matrix:
-      py37_conda_mkl:
-        PYTHON_VERSION: '3.7'
-        CHECK_WARNINGS: 'true'
-        PYTHON_ARCH: '64'
-        PYTEST_VERSION: '*'
-        COVERAGE: 'true'
-      py36_pip_openblas_32bit:
-        PYTHON_VERSION: '3.6'
-        PYTHON_ARCH: '32'
+# - template: build_tools/azure/windows.yml
+#   parameters:
+#     name: Windows
+#     vmImage: vs2017-win2016
+#     dependsOn: [linting]
+#     condition: and(ne(variables['Build.Reason'], 'Schedule'), succeeded('linting'))
+#     matrix:
+#       py37_conda_mkl:
+#         PYTHON_VERSION: '3.7'
+#         CHECK_WARNINGS: 'true'
+#         PYTHON_ARCH: '64'
+#         PYTEST_VERSION: '*'
+#         COVERAGE: 'true'
+#       py36_pip_openblas_32bit:
+#         PYTHON_VERSION: '3.6'
+#         PYTHON_ARCH: '32'
diff --git a/build_tools/azure/install.sh b/build_tools/azure/install.sh
@@ -94,7 +94,7 @@ elif [[ "$DISTRIB" == "conda-pip-latest" ]]; then
     python -m pip install -U pip
     python -m pip install pytest==$PYTEST_VERSION pytest-cov
 
-    python -m pip install pandas matplotlib pyamg scikit-image
+    python -m pip install pandas matplotlib pyamg scikit-image jax jaxlib
     # do not install dependencies for lightgbm since it requires scikit-learn
     python -m pip install lightgbm --no-deps
 elif [[ "$DISTRIB" == "conda-pip-scipy-dev" ]]; then

diff --git a/build_tools/azure/test_script.sh b/build_tools/azure/test_script.sh
@@ -42,5 +42,5 @@ cp setup.cfg $TEST_DIR
 cd $TEST_DIR
 
 set -x
-$TEST_CMD --pyargs sklearn
+$TEST_CMD --pyargs sklearn.decomposition
 set +x
diff --git a/sklearn/_config.py b/sklearn/_config.py
@@ -8,6 +8,7 @@
     'working_memory': int(os.environ.get('SKLEARN_WORKING_MEMORY', 1024)),
     'print_changed_only': True,
     'display': 'text',
+    'enable_duck_array': False,
 }
 
 
@@ -28,7 +29,8 @@ def get_config():
 
 
 def set_config(assume_finite=None, working_memory=None,
-               print_changed_only=None, display=None):
+               print_changed_only=None, display=None,
+               enable_duck_array=None):
     """Set global scikit-learn configuration
 
     .. versionadded:: 0.19
@@ -80,6 +82,8 @@ def set_config(assume_finite=None, working_memory=None,
         _global_config['print_changed_only'] = print_changed_only
     if display is not None:
         _global_config['display'] = display
+    if enable_duck_array is not None:
+        _global_config['enable_duck_array'] = enable_duck_array
 
 
 @contextmanager

diff --git a/sklearn/decomposition/_pca.py b/sklearn/decomposition/_pca.py
@@ -14,7 +14,6 @@
 import numbers
 
 import numpy as np
-from scipy import linalg
 from scipy.special import gammaln
 from scipy.sparse import issparse
 from scipy.sparse.linalg import svds
@@ -26,6 +25,7 @@
 from ..utils.extmath import stable_cumsum
 from ..utils.validation import check_is_fitted
 from ..utils.validation import _deprecate_positional_args
+from ..utils import _get_array_module
 
 
 def _assess_dimension(spectrum, rank, n_samples):
@@ -396,6 +396,7 @@ def _fit(self, X):
 
         X = self._validate_data(X, dtype=[np.float64, np.float32],
                                 ensure_2d=True, copy=self.copy)
+        npx = _get_array_module(X)
 
         # Handle n_components==None
         if self.n_components is None:
@@ -420,14 +421,14 @@ def _fit(self, X):
 
         # Call different fits for either full or truncated SVD
         if self._fit_svd_solver == 'full':
-            return self._fit_full(X, n_components)
+            return self._fit_full(X, n_components, npx=npx)
         elif self._fit_svd_solver in ['arpack', 'randomized']:
             return self._fit_truncated(X, n_components, self._fit_svd_solver)
         else:
             raise ValueError("Unrecognized svd_solver='{0}'"
                              "".format(self._fit_svd_solver))
 
-    def _fit_full(self, X, n_components):
+    def _fit_full(self, X, n_components, npx=np):
         """Fit the model by computing full SVD on X"""
         n_samples, n_features = X.shape
 
@@ -448,12 +449,12 @@ def _fit_full(self, X, n_components):
                                  % (n_components, type(n_components)))
 
         # Center data
-        self.mean_ = np.mean(X, axis=0)
+        self.mean_ = npx.mean(X, axis=0)
         X -= self.mean_
 
-        U, S, Vt = linalg.svd(X, full_matrices=False)
+        U, S, Vt = npx.linalg.svd(X, full_matrices=False)
         # flip eigenvectors' sign to enforce deterministic output
-        U, Vt = svd_flip(U, Vt)
+        U, Vt = svd_flip(U, Vt, npx=npx)
 
         components_ = Vt
 

diff --git a/sklearn/decomposition/tests/test_pca.py b/sklearn/decomposition/tests/test_pca.py
@@ -3,6 +3,7 @@
 
 import pytest
 
+import sklearn
 from sklearn.utils._testing import assert_allclose
 
 from sklearn import datasets
@@ -638,3 +639,30 @@ def test_assess_dimesion_rank_one():
     assert np.isfinite(_assess_dimension(s, rank=1, n_samples=n_samples))
     for rank in range(2, n_features):
         assert _assess_dimension(s, rank, n_samples) == -np.inf
+
+
+# XXX: it should be possible to support 'randomized' by adding npx=np
+# in appropriate locations. The 'arpack' svd_solver, on the other-hand,
+# cannot easily be adapted to work on non-numpy allocated arrays.
+# @pytest.mark.parametrize('svd_solver', ["full", "randomized", "auto"])
+@pytest.mark.parametrize('svd_solver', ["full"])
+@pytest.mark.parametrize('copy', [True, False])
+def test_pca_jax_data(svd_solver, copy):
+    jnp = pytest.importorskip("jax.numpy")
+    X_np = np.random.RandomState(42).randn(1000, 100)
+    X_np = X_np.astype(np.float32)
+    X_jnp = jnp.asarray(X_np)
+
+    pca_np = PCA(n_components=3, svd_solver=svd_solver, copy=copy,
+                 random_state=0)
+    X_pca_np = pca_np.fit_transform(X_np)
+
+    with sklearn.config_context(enable_duck_array=True):
+        pca_jnp = PCA(**pca_np.get_params())
+        X_pca_jnp = pca_jnp.fit_transform(X_jnp)
+
+    assert isinstance(X_pca_jnp, type(X_jnp))
+    assert isinstance(pca_jnp.components_, type(X_jnp))
+
+    assert_allclose(X_pca_np, X_pca_jnp, atol=1e-3)
+    assert_allclose(pca_np.components_, pca_jnp.components_, atol=1e-3)