Repository: scikit-learn/scikit-learn
File: sklearn/preprocessing/tests/test_label.py (Python)

Code before:
        [
            [1, 0, 2, 9],
            [0, 1, 2, 9],
            [[0, 1, 0, 0], [1, 0, 0, 0], [0, 0, 1, 0], [0, 0, 0, 1]],
        ],
    ],
)
@pytest.mark.parametrize(
    "array_namespace, device_, dtype_name", yield_namespace_device_dtype_combinations()
)
def test_label_binarizer_array_api_compliance(
    y, classes, expected, array_namespace, device_, dtype_name
):
    """Test that :class:`LabelBinarizer` works correctly with the Array API for binary
    and multi-class inputs for numerical labels and non-sparse outputs.
    """
    xp = _array_api_for_tests(array_namespace, device_)

    y_np = np.asarray(y)

    with config_context(array_api_dispatch=True):
        y = xp.asarray(y, device=device_)

        # `sparse_output=True` is not allowed for non-NumPy namespace.
        # Similarly, if `LabelBinarizer` is fitted on sparse matrix,
        # then inverse-transforming non-NumPy array is not allowed.
        if not _is_numpy_namespace(xp):
            with pytest.raises(ValueError):
                LabelBinarizer(sparse_output=True).fit(y)

            with pytest.raises(ValueError):
                lb_np = LabelBinarizer(sparse_output=True).fit(y_np)
                lb_np.transform(y)

            with pytest.raises(ValueError):
                lb_sparse = LabelBinarizer().fit(y_np)
                lb_sparse.sparse_input_ = True
                lb_sparse.inverse_transform(xp.asarray(expected, device=device_))

        # Shouldn't raise error in both `fit` and `transform` when `sparse_output=False`
        lb_xp = LabelBinarizer()

        binarized = lb_xp.fit_transform(y)
        assert get_namespace(binarized)[0].__name__ == xp.__name__
        assert "int" in str(binarized.dtype)
        assert device(binarized) == device(y)
        assert_array_equal(_convert_to_numpy(binarized, xp=xp), np.asarray(expected))

        fitted_classes = lb_xp.classes_
        assert get_namespace(fitted_classes)[0].__name__ == xp.__name__
        assert device(fitted_classes) == device(y)

Review comment (suggestion):
```suggestion
        # `sparse_output=True` is not allowed for non-NumPy namespaces.
        # Similarly, if `LabelBinarizer` is fitted on a sparse matrix,
        # then inverse-transforming non-NumPy arrays is not allowed.
```

Code after:
        [
            [1, 0, 2, 9],
            [0, 1, 2, 9],
            [[0, 1, 0, 0], [1, 0, 0, 0], [0, 0, 1, 0], [0, 0, 0, 1]],
        ],
    ],
)
@pytest.mark.parametrize(
    "array_namespace, device_, dtype_name", yield_namespace_device_dtype_combinations()
)
def test_label_binarizer_array_api_compliance(
    y, classes, expected, array_namespace, device_, dtype_name
):
    """Test that :class:`LabelBinarizer` works correctly with the Array API for binary
    and multi-class inputs for numerical labels and non-sparse outputs.
    """
    xp = _array_api_for_tests(array_namespace, device_)

    y_np = np.asarray(y)

    with config_context(array_api_dispatch=True):
        y = xp.asarray(y, device=device_)

        # `sparse_output=True` is not allowed for non-NumPy namespaces.
        # Similarly, if `LabelBinarizer` is fitted on a sparse matrix,
        # then inverse-transforming non-NumPy arrays is not allowed.
        if not _is_numpy_namespace(xp):
            sparse_output_msg = "`sparse_output=True` is not supported for array API"

            with pytest.raises(ValueError, match=sparse_output_msg):
                LabelBinarizer(sparse_output=True).fit(y)

            lb_np = LabelBinarizer(sparse_output=True).fit(y_np)
            with pytest.raises(ValueError, match=sparse_output_msg):
                lb_np.transform(y)

            lb_sparse = LabelBinarizer().fit(y_np)
            lb_sparse.sparse_input_ = True
            sparse_input_msg = (
                "`LabelBinarizer` was fitted on a sparse matrix, and therefore cannot"
            )
            with pytest.raises(ValueError, match=sparse_input_msg):
                lb_sparse.inverse_transform(xp.asarray(expected, device=device_))

        # Shouldn't raise error in both `fit` and `transform` when `sparse_output=False`
        lb_xp = LabelBinarizer()

        binarized = lb_xp.fit_transform(y)
        assert get_namespace(binarized)[0].__name__ == xp.__name__
        assert "int" in str(binarized.dtype)
        assert device(binarized) == device(y)