scikit-learn · jeromedockes · Mar 13, 2024 · Mar 13, 2024 · Mar 14, 2024 · Mar 20, 2024
diff --git a/doc/modules/array_api.rst b/doc/modules/array_api.rst
@@ -101,6 +101,7 @@ Estimators
 - :class:`preprocessing.MaxAbsScaler`
 - :class:`preprocessing.MinMaxScaler`
 - :class:`preprocessing.Normalizer`
+- :class:`preprocessing.LabelBinarizer`
 
 Metrics
 -------

diff --git a/doc/whats_new/v1.6.rst b/doc/whats_new/v1.6.rst
@@ -38,7 +38,9 @@ See :ref:`array_api` for more details.
 
 **Classes:**
 
--
+- :class:`sklearn.preprocessing.LabelBinarizer` now supports Array API compliant inputs.
+  :pr:`28626` by :user:`Jérôme Dockès <jeromedockes>` and :user:`Olivier Grisel <ogrisel>`.
+
 
 Metadata Routing
 ----------------

diff --git a/sklearn/preprocessing/_label.py b/sklearn/preprocessing/_label.py
@@ -17,6 +17,7 @@
 
 from ..base import BaseEstimator, TransformerMixin, _fit_context
 from ..utils import column_or_1d
+from ..utils._array_api import _convert_to_numpy, device, get_namespace
 from ..utils._encode import _encode, _unique
 from ..utils._param_validation import Interval, validate_params
 from ..utils.multiclass import type_of_target, unique_labels
@@ -303,7 +304,8 @@ def fit(self, y):
             raise ValueError("y has 0 samples: %r" % y)
 
         self.sparse_input_ = sp.issparse(y)
-        self.classes_ = unique_labels(y)
+        xp, _ = get_namespace(y)
+        self.classes_ = _convert_to_numpy(unique_labels(y), xp)
         return self
 
     def fit_transform(self, y):
@@ -396,6 +398,21 @@ def inverse_transform(self, Y, threshold=None):
         """
         check_is_fitted(self)
 
+        # LabelBinarizer supports Array API compatibility for convenience when
+        # used as a sub-component of an classifier that does. However
+        # label_binarize internally uses a NumPy copy of the data because
+        # all the operations are meant to construct the backing NumPy arrays of a
+        # scipy.sparse CSR datastructure even when sparse_output=False.
+        #
+        # In the future, we might consider a dedicated code path for the
+        # sparse_output=False case that would directly be implemented using Array
+        # API without the intermediate NumPy conversion and scipy.sparse
+        # datastructure.
+        xp, is_array_api_compliant = get_namespace(Y)
+        device_ = device(Y) if is_array_api_compliant else None
+        if not sp.issparse(Y):
+            Y = _convert_to_numpy(Y, xp)
+
         if threshold is None:
             threshold = (self.pos_label + self.neg_label) / 2.0
 
@@ -410,11 +427,13 @@ def inverse_transform(self, Y, threshold=None):
             y_inv = sp.csr_matrix(y_inv)
         elif sp.issparse(y_inv):
             y_inv = y_inv.toarray()
+        if is_array_api_compliant and not sp.issparse(y_inv):
+            y_inv = xp.asarray(y_inv, device=device_)
 
         return y_inv
 
     def _more_tags(self):
-        return {"X_types": ["1dlabels"]}
+        return {"X_types": ["1dlabels"], "array_api_support": True}
 
 
 @validate_params(
@@ -487,6 +506,18 @@ def label_binarize(y, *, classes, neg_label=0, pos_label=1, sparse_output=False)
            [0],
            [1]])
     """
+    # label_binarize supports Array API compatibility for convenience when
+    # LabelBinarizer is used as a sub-component of an classifier that does.
+    # However label_binarize internally uses a NumPy copy of the data because
+    # all the operations are meant to construct the backing NumPy arrays of a
+    # scipy.sparse CSR datastructure even when sparse_output=False.
+    y_xp, y_is_array_api = get_namespace(y)
+    if y_is_array_api:
+        device_ = device(y)
+        y = _convert_to_numpy(y, y_xp)
+    classes_xp, classes_is_array_api = get_namespace(classes)
+    if classes_is_array_api:
+        classes = _convert_to_numpy(classes, classes_xp)
     if not isinstance(y, list):
         # XXX Workaround that will be removed when list of list format is
         # dropped
@@ -535,7 +566,10 @@ def label_binarize(y, *, classes, neg_label=0, pos_label=1, sparse_output=False)
             else:
                 Y = np.zeros((len(y), 1), dtype=int)
                Y += neg_label
-                return Y
+
+                if not y_is_array_api:
+                    return Y
+                return y_xp.asarray(Y, device=device_)
         elif len(classes) >= 3:
             y_type = "multiclass"
 
@@ -595,7 +629,9 @@ def label_binarize(y, *, classes, neg_label=0, pos_label=1, sparse_output=False)
         else:
             Y = Y[:, -1].reshape((-1, 1))
 
-    return Y
+    if not y_is_array_api:
+        return Y
+    return y_xp.asarray(Y, devic
A92E
e=device_)
 
 
 def _inverse_binarize_multiclass(y, classes):

diff --git a/sklearn/preprocessing/tests/test_label.py b/sklearn/preprocessing/tests/test_label.py
@@ -2,7 +2,7 @@
 import pytest
 from scipy.sparse import issparse
 
-from sklearn import datasets
+from sklearn import config_context, datasets
 from sklearn.preprocessing._label import (
     LabelBinarizer,
     LabelEncoder,
@@ -11,7 +11,16 @@
     _inverse_binarize_thresholding,
     label_binarize,
 )
-from sklearn.utils._testing import assert_array_equal, ignore_warnings
+from sklearn.utils._array_api import (
+    _convert_to_numpy,
+    get_namespace,
+    yield_namespace_device_dtype_combinations,
+)
+from sklearn.utils._testing import (
+    _array_api_for_tests,
+    assert_array_equal,
+    ignore_warnings,
+)
 from sklearn.utils.fixes import (
     COO_CONTAINERS,
     CSC_CONTAINERS,
@@ -216,6 +225,35 @@ def test_label_binarizer_sparse_errors(csr_container):
         )
 
 
+@pytest.mark.parametrize(
+    "array_namespace, device, dtype_name", yield_namespace_device_dtype_combinations()
+)
+@pytest.mark.parametrize(
+    "y",
+    [
+        np.array([1, 0, 2]),
+        np.array([0, 0, 0]),
+        np.array([1, 0, 0]),
+        np.array([[0, 1, 1], [1, 0, 1]]),
+    ],
+)
+def test_label_binarizer_array_api(y, array_namespace, device, dtype_name):
+    xp = _array_api_for_tests(array_namespace, device)
+    xp_y = xp.asarray(y, device=device)
+    xp_lb = LabelBinarizer(sparse_output=False)
+    with config_context(array_api_dispatch=True):
+        xp_fit_transformed = xp_lb.fit_transform(xp_y)
+        xp_transformed = xp_lb.transform(xp_y)
+        xp_inv_transformed = xp_lb.inverse_transform(xp_fit_transformed)
+        np_lb = LabelBinarizer(sparse_output=False)
+        np_transformed = np_lb.fit_transform(y)
+        assert get_namespace(xp_fit_transformed)[0].__name__ == xp.__name__
+        assert get_namespace(xp_transformed)[0].__name__ == xp.__name__
+        assert get_namespace(xp_inv_transformed)[0].__name__ == xp.__name__
+        assert_array_equal(_convert_to_numpy(xp_fit_transformed, xp), np_transformed)
+        assert_array_equal(_convert_to_numpy(xp_inv_transformed, xp), y)
+
+
 @pytest.mark.parametrize(
     "values, classes, unknown",
     [