Repository: scikit-learn/scikit-learn
File: sklearn/utils/tests/test_array_api.py (Python)

Code before:
)
def test_add_to_diagonal(array_namespace, device_, dtype_name):
    """Check `_add_to_diagonal` consistent between array API xp and numpy namespace."""
    xp = _array_api_for_tests(array_namespace, device_)
    np_xp = _array_api_for_tests("numpy", None)

    array_np = numpy.zeros((3, 4), dtype=dtype_name)
    array_xp = xp.asarray(array_np.copy(), device=device_)

    add_val = [1, 2, 3]
    _fill_diagonal(array_np, value=add_val, xp=np_xp)
    with config_context(array_api_dispatch=True):
        _fill_diagonal(array_xp, value=add_val, xp=xp)

    assert_array_equal(_convert_to_numpy(array_xp, xp=xp), array_np)


@pytest.mark.parametrize("csr_container", CSR_CONTAINERS)
@pytest.mark.parametrize("dispatch", [True, False])
def test_sparse_device(csr_container, dispatch):
    a, b = csr_container(numpy.array([[1]])), csr_container(numpy.array([[2]]))
    if dispatch and os.environ.get("SCIPY_ARRAY_API") is None:
        raise SkipTest("SCIPY_ARRAY_API is not set: not checking array_api input")
    with config_context(array_api_dispatch=dispatch):
        assert device(a, b) is None
        assert device(a, numpy.array([1])) == ("cpu" if dispatch else None)
        assert get_namespace_and_device(a, b)[2] is None
        assert get_namespace_and_device(a, numpy.array([1]))[2] == (
            "cpu" if dispatch else None
        )


@pytest.mark.parametrize(
    "namespace, device, dtype_name",
    yield_namespace_device_dtype_combinations(),
    ids=_get_namespace_device_dtype_ids,
)
@pytest.mark.parametrize("axis", [None, 0, 1])
def test_median(namespace, device, dtype_name, axis):
    # Note: depending on the value of `axis`, this test will compare median
    # computations on arrays of even (4) or odd (5) numbers of elements, hence
    # will test for median computation with and without interpolation to check
    # that array API namespaces yield consistent results even when the median is
    # not mathematically uniquely defined.
    xp = _array_api_for_tests(namespace, device)
    rng = numpy.random.RandomState(0)

    X_np = rng.uniform(low=0.0, high=1.0, size=(5, 4)).astype(dtype_name)
    result_np = numpy.median(X_np, axis=axis)

    X_xp = xp.asarray(X_np, device=device)

Review comment (bug):
That will fail for numpy<2, right? In this case the numpy array has no `.device` so even when `dispatch is True`, the device will be `None`.

To be honest, the simplest may be to decide you need numpy>=2 for array API.

Right now we have scipy 1.14.0 as the minimum version (released 2024-06-24) for array API and numpy 2.0 was released around the same time (2024-06-16).

Code after:
)
def test_add_to_diagonal(array_namespace, device_, dtype_name):
    """Check `_add_to_diagonal` consistent between array API xp and numpy namespace."""
    xp = _array_api_for_tests(array_namespace, device_)
    np_xp = _array_api_for_tests("numpy", None)

    array_np = numpy.zeros((3, 4), dtype=dtype_name)
    array_xp = xp.asarray(array_np.copy(), device=device_)

    add_val = [1, 2, 3]
    _fill_diagonal(array_np, value=add_val, xp=np_xp)
    with config_context(array_api_dispatch=True):
        _fill_diagonal(array_xp, value=add_val, xp=xp)

    assert_array_equal(_convert_to_numpy(array_xp, xp=xp), array_np)


@pytest.mark.parametrize("csr_container", CSR_CONTAINERS)
@pytest.mark.parametrize("dispatch", [True, False])
def test_sparse_device(csr_container, dispatch):
    np_arr = numpy.array([1])
    # For numpy < 2, the device attribute is not available on numpy arrays
    expected_numpy_array_device = getattr(np_arr, "device", None) if dispatch else None
    a, b = csr_container(numpy.array([[1]])), csr_container(numpy.array([[2]]))
    if dispatch and os.environ.get("SCIPY_ARRAY_API") is None:
        raise SkipTest("SCIPY_ARRAY_API is not set: not checking array_api input")
    with config_context(array_api_dispatch=dispatch):
        assert device(a, b) is None
        assert device(a, np_arr) == expected_numpy_array_device
        assert get_namespace_and_device(a, b)[2] is None
        assert get_namespace_and_device(a, np_arr)[2] == expected_numpy_array_device


@pytest.mark.parametrize(
    "namespace, device, dtype_name",
    yield_namespace_device_dtype_combinations(),
    ids=_get_namespace_device_dtype_ids,
)
@pytest.mark.parametrize("axis", [None, 0, 1])
def test_median(namespace, device, dtype_name, axis):
    # Note: depending on the value of `axis`, this test will compare median
    # computations on arrays of even (4) or odd (5) numbers of elements, hence
    # will test for median computation with and without interpolation to check
    # that array API namespaces yield consistent results even when the median is
    # not mathematically uniquely defined.
    xp = _array_api_for_tests(namespace, device)
    rng = numpy.random.RandomState(0)

    X_np = rng.uniform(low=0.0, high=1.0, size=(5, 4)).astype(dtype_name)
    result_np = numpy.median(X_np, axis=axis)