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

Code before:
@pytest.mark.parametrize(
    "array_namespace, device_, dtype_name", yield_namespace_device_dtype_combinations()
)
def test_label_binarize_array_api_compliance(
    y, classes, expected, array_namespace, device_, dtype_name
):
    """Test that :func:`label_binarize` 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_)
    xp_is_numpy = _is_numpy_namespace(xp)
    numeric_dtype = np.issubdtype(np.asarray(y).dtype, np.integer) and np.issubdtype(
        np.asarray(classes).dtype, np.integer
    )

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

        if numeric_dtype:
            # `sparse_output=True` is not allowed for non-NumPy namespace
            if not xp_is_numpy:
                msg = "`sparse_output=True` is not supported for array API "
                with pytest.raises(ValueError, match=msg):
                    label_binarize(y=y, classes=classes, sparse_output=True)

            # Numeric class labels should not raise any errors for non-NumPy namespace
            binarized = label_binarize(y, classes=classes)
            expected = np.asarray(expected, dtype=int)

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

        if not xp_is_numpy and not numeric_dtype:
            msg = "`classes` contains unsupported dtype for array API "
            with pytest.raises(ValueError, match=msg):
                label_binarize(y=y, classes=classes)


@pytest.mark.parametrize("csr_container", CSR_CONTAINERS)
def test_inverse_binarize_multiclass(csr_container):
    got = _inverse_binarize_multiclass(
        csr_container([[0, 1, 0], [-1, 0, -1], [0, 0, 0]]), np.arange(3)
    )
    assert_array_equal(got, np.array([1, 1, 0]))


def test_nan_label_encoder():
    """Check that label encoder encodes nans in transform.


Review comment (suggestion):
```suggestion
            # Numeric class labels should not raise any errors for non-NumPy namespaces
            # when `sparse_output=False`
```

Code after:
@pytest.mark.parametrize(
    "array_namespace, device_, dtype_name", yield_namespace_device_dtype_combinations()
)
def test_label_binarize_array_api_compliance(
    y, classes, expected, array_namespace, device_, dtype_name
):
    """Test that :func:`label_binarize` 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_)
    xp_is_numpy = _is_numpy_namespace(xp)
    numeric_dtype = np.issubdtype(np.asarray(y).dtype, np.integer) and np.issubdtype(
        np.asarray(classes).dtype, np.integer
    )

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

        if numeric_dtype:
            # `sparse_output=True` is not allowed for non-NumPy namespaces
            if not xp_is_numpy:
                msg = "`sparse_output=True` is not supported for array API "
                with pytest.raises(ValueError, match=msg):
                    label_binarize(y=y, classes=classes, sparse_output=True)

            # Numeric class labels should not raise any errors for non-NumPy namespaces
            binarized = label_binarize(y, classes=classes)
            expected = np.asarray(expected, dtype=int)

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

        if not xp_is_numpy and not numeric_dtype:
            msg = "`classes` contains unsupported dtype for array API "
            with pytest.raises(ValueError, match=msg):
                label_binarize(y=y, classes=classes)


@pytest.mark.parametrize("csr_container", CSR_CONTAINERS)
def test_inverse_binarize_multiclass(csr_container):
    got = _inverse_binarize_multiclass(
        csr_container([[0, 1, 0], [-1, 0, -1], [0, 0, 0]]), np.arange(3)
    )
    assert_array_equal(got, np.array([1, 1, 0]))


def test_nan_label_encoder():
    """Check that label encoder encodes nans in transform.