scikit-learn · ogrisel · Jul 1, 2025 · Jun 17, 2025 · Jun 18, 2025 · Jun 18, 2025
diff --git a/doc/modules/array_api.rst b/doc/modules/array_api.rst
@@ -117,6 +117,7 @@ Estimators
 - :class:`preprocessing.MaxAbsScaler`
 - :class:`preprocessing.MinMaxScaler`
 - :class:`preprocessing.Normalizer`
+- :class:`preprocessing.PolynomialFeatures`
 - :class:`mixture.GaussianMixture` (with `init_params="random"` or
   `init_params="random_from_data"` and `warm_start=False`)
 

diff --git a/doc/whats_new/upcoming_changes/array-api/31580.feature.rst b/doc/whats_new/upcoming_changes/array-api/31580.feature.rst
@@ -0,0 +1,2 @@
+- :class:`preprocessing.PolynomialFeatures` now supports array API compatible inputs.
+  By :user:`Omar Salman <OmarManzoor>`
diff --git a/sklearn/preprocessing/_polynomial.py b/sklearn/preprocessing/_polynomial.py
@@ -15,6 +15,12 @@
 from scipy.interpolate import BSpline
 from scipy.special import comb
 
+from sklearn.utils._array_api import (
+    _is_numpy_namespace,
+    get_namespace_and_device,
+    supported_float_dtypes,
+)
+
 from ..base import BaseEstimator, TransformerMixin, _fit_context
 from ..utils import check_array
 from ..utils._mask import _get_mask
@@ -416,18 +422,18 @@ def transform(self, X):
             `csr_matrix`.
         """
         check_is_fitted(self)
-
+        xp, _, device_ = get_namespace_and_device(X)
         X = validate_data(
             self,
             X,
             order="F",
-            dtype=FLOAT_DTYPES,
+            dtype=supported_float_dtypes(xp=xp, device=device_),
             reset=False,
             accept_sparse=("csr", "csc"),
         )
 
         n_samples, n_features = X.shape
-        max_int32 = np.iinfo(np.int32).max
+        max_int32 = xp.iinfo(xp.int32).max
         if sparse.issparse(X) and X.format == "csr":
             if self._max_degree > 3:
                 return self.transform(X.tocsc()).tocsr()
@@ -497,8 +503,19 @@ def transform(self, X):
         else:
             # Do as if _min_degree = 0 and cut down array after the
             # computation, i.e. use _n_out_full instead of n_output_features_.
-            XP = np.empty(
-                shape=(n_samples, self._n_out_full), dtype=X.dtype, order=self.order
+            order_kwargs = {}
+            if _is_numpy_namespace(xp=xp):
+                order_kwargs["order"] = self.order
+            elif self.order == "F":
+                raise ValueError(
+                    "PolynomialFeatures does not support order='F' for non-numpy arrays"
+                )
+
+            XP = xp.empty(
+                shape=(n_samples, self._n_out_full),
+                dtype=X.dtype,
+                device=device_,
+                **order_kwargs,
             )
 
             # What follows is a faster implementation of:
@@ -544,12 +561,18 @@ def transform(self, X):
                         break
                     # XP[:, start:end] are terms of degree d - 1
                     # that exclude feature #feature_idx.
-                    np.multiply(
-                        XP[:, start:end],
-                        X[:, feature_idx : feature_idx + 1],
-                        out=XP[:, current_col:next_col],
-                        casting="no",
-                    )
+                    if _is_numpy_namespace(xp):
+                        # numpy performs this multiplication in place
+                        np.multiply(
+                            XP[:, start:end],
+                            X[:, feature_idx : feature_idx + 1],
+                            out=XP[:, current_col:next_col],
+                            casting="no",
+                        )
+                    else:
+                        XP[:, current_col:next_col] = xp.multiply(
+                            XP[:, start:end], X[:, feature_idx : feature_idx + 1]
+                        )
                     current_col = next_col
 
                 new_index.append(current_col)
@@ -558,19 +581,23 @@ def transform(self, X):
             if self._min_degree > 1:
                 n_XP, n_Xout = self._n_out_full, self.n_output_features_
                 if self.include_bias:
-                    Xout = np.empty(
-                        shape=(n_samples, n_Xout), dtype=XP.dtype, order=self.order
+                    Xout = xp.empty(
+                        shape=(n_samples, n_Xout),
+                        dtype=XP.dtype,
+                        device=device_,
+                        **order_kwargs,
                     )
                     Xout[:, 0] = 1
                     Xout[:, 1:] = XP[:, n_XP - n_Xout + 1 :]
                 else:
-                    Xout = XP[:, n_XP - n_Xout :].copy()
+                    Xout = xp.asarray(XP[:, n_XP - n_Xout :], copy=True)
                 XP = Xout
         return XP
 
     def __sklearn_tags__(self):
         tags = super().__sklearn_tags__()
         tags.input_tags.sparse = True
+        tags.array_api_support = True
         return tags
 
 

diff --git a/sklearn/preprocessing/tests/test_polynomial.py b/sklearn/preprocessing/tests/test_polynomial.py
@@ -8,6 +8,7 @@
 from scipy.interpolate import BSpline
 from scipy.sparse import random as sparse_random
 
+from sklearn._config import config_context
 from sklearn.linear_model import LinearRegression
 from sklearn.pipeline import Pipeline
 from sklearn.preprocessing import (
@@ -18,8 +19,17 @@
 from sklearn.preprocessing._csr_polynomial_expansion import (
     _get_sizeof_LARGEST_INT_t,
 )
+from sklearn.utils._array_api import (
+    _convert_to_numpy,
+    _get_namespace_device_dtype_ids,
+    _is_numpy_namespace,
+    device,
+    get_namespace,
+    yield_namespace_device_dtype_combinations,
+)
 from sklearn.utils._mask import _get_mask
 from sklearn.utils._testing import (
+    _array_api_for_tests,
     assert_allclose_dense_sparse,
     assert_array_almost_equal,
 )
@@ -1336,3 +1346,64 @@ def test_csr_polynomial_expansion_windows_fail(csr_container):
         X_trans = pf.fit_transform(X)
         for idx in range(3):
             assert X_trans[0, expected_indices[idx]] == pytest.approx(1.0)
+
+
+@pytest.mark.parametrize(
+    "array_namespace, device_, dtype_name",
+    yield_namespace_device_dtype_combinations(),
+    ids=_get_namespace_device_dtype_ids,
+)
+@pytest.mark.parametrize("interaction_only", [True, False])
+@pytest.mark.parametrize("include_bias", [True, False])
+@pytest.mark.parametrize("degree", [2, (2, 2), 3, (3, 3)])
+def test_polynomial_features_array_api_compliance(
+    two_features_degree3,
+    degree,
+    include_bias,
+    interaction_only,
+    array_namespace,
+    device_,
+    dtype_name,
+):
+    """Test array API compliance for PolynomialFeatures on 2 features up to degree 3."""
+    xp = _array_api_for_tests(array_namespace, device_)
+    X, _ = two_features_degree3
+    X_np = X.astype(dtype_name)
+    X_xp = xp.asarray(X_np, device=device_)
+    with config_context(array_api_dispatch=True):
+        tf_np = PolynomialFeatures(
+            degree=degree, include_bias=include_bias, interaction_only=interaction_only
+        ).fit(X_np)
+
+        tf_xp = PolynomialFeatures(
+            degree=degree, include_bias=include_bias, interaction_only=interaction_only
+        ).fit(X_xp)
+        out_np = tf_np.transform(X_np)
+        out_xp = tf_xp.transform(X_xp)
+        assert_allclose(_convert_to_numpy(out_xp, xp=xp), out_np)
+        assert get_namespace(out_xp)[0].__name__ == xp.__name__
+        assert device(out_xp) == device(X_xp)
+        assert out_xp.dtype == X_xp.dtype
+
+
+@pytest.mark.parametrize(
+    "array_namespace, device_, dtype_name",
+    yield_namespace_device_dtype_combinations(),
+    ids=_get_namespace_device_dtype_ids,
+)
+def test_polynomial_features_array_api_raises_on_order_F(
+    array_namespace, device_, dtype_name
+):
+    """Test that PolynomialFeatures with order='F' raises ValueError on
+    array API namespaces other than numpy."""
+    xp = _array_api_for_tests(array_namespace, device_)
+    X = np.arange(6).reshape((3, 2)).astype(dtype_name)
+    X_xp = xp.asarray(X, device=device_)
+    msg = "PolynomialFeatures does not support order='F' for non-numpy arrays"
+    with config_context(array_api_dispatch=True):
+        pf = PolynomialFeatures(order="F").fit(X_xp)
+        if _is_numpy_namespace(xp):  # Numpy should not raise
+            pf.transform(X_xp)
+        else:
+            with pytest.raises(ValueError, match=msg):
+                pf.transform(X_xp)
diff --git a/sklearn/utils/_array_api.py b/sklearn/utils/_array_api.py
@@ -290,7 +290,7 @@ def _isdtype_single(dtype, kind, *, xp):
         return dtype == kind
 
 
-def supported_float_dtypes(xp):
+def supported_float_dtypes(xp, device=None):
     """Supported floating point types for the namespace.
 
     Note: float16 is not officially part of the Array API spec at the
@@ -299,10 +299,18 @@ def supported_float_dtypes(xp):
 
     https://data-apis.org/array-api/latest/API_specification/data_types.html
     """
+    dtypes_dict = xp.__array_namespace_info__().dtypes(
+        kind="real floating", device=device
+    )
+    valid_float_dtypes = []
+    for dtype_key in ("float64", "float32"):
+        if dtype_key in dtypes_dict:
+            valid_float_dtypes.append(dtypes_dict[dtype_key])
+
     if hasattr(xp, "float16"):
-        return (xp.float64, xp.float32, xp.float16)
-    else:
-        return (xp.float64, xp.float32)
+        valid_float_dtypes.append(xp.float16)
+
+    return tuple(valid_float_dtypes)
 
 
 def ensure_common_namespace_device(reference, *arrays):

diff --git a/sklearn/utils/tests/test_array_api.py b/sklearn/utils/tests/test_array_api.py
@@ -33,6 +33,7 @@
     get_namespace_and_device,
     indexing_dtype,
     np_compat,
+    supported_float_dtypes,
     yield_namespace_device_dtype_combinations,
 )
 from sklearn.utils._testing import (
@@ -777,3 +778,20 @@ def test_logsumexp_like_scipy_logsumexp(array_namespace, device_, dtype_name, ax
         res_xp_2 = _logsumexp(array_xp_2, axis=axis)
         res_xp_2 = _convert_to_numpy(res_xp_2, xp)
         assert_allclose(res_np_2, res_xp_2, rtol=rtol)
+
+
+@pytest.mark.parametrize(
+    ("namespace", "device_", "expected_types"),
+    [
+        ("numpy", None, ("float64", "float32", "float16")),
+        ("array_api_strict", None, ("float64", "float32")),
+        ("torch", "cpu", ("float64", "float32", "float16")),
+        ("torch", "cuda", ("float64", "float32", "float16")),
+        ("torch", "mps", ("float32", "float16")),
+    ],
+)
+def test_supported_float_types(namespace, device_, expected_types):
+    xp = _array_api_for_tests(namespace, device_)
+    float_types = supported_float_dtypes(xp, device=device_)
+    expected = tuple(getattr(xp, dtype_name) for dtype_name in expected_types)
+    assert float_types == expected
Original file line number	Diff line number	Diff line change
		@@ -0,0 +1,2 @@
		- :class:`preprocessing.PolynomialFeatures` now supports array API compatible inputs.
		By :user:`Omar Salman <OmarManzoor>`