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

Code before:
    with config_context(array_api_dispatch=True):
        _fill_or_add_to_diagonal(array_xp, value=1, xp=xp, add_value=False, wrap=wrap)

    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])) is None
        assert get_namespace_and_device(a, b)[2] is None
        assert get_namespace_and_device(a, numpy.array([1]))[2] is 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):
    xp = _array_api_for_tests(namespace, device)
    rng = numpy.random.RandomState(0)

    X_np = numpy.array(rng.random_sample((5, 4)), dtype=dtype_name)
    result_np = numpy.median(X_np, axis=axis)

    X_xp = xp.asarray(X_np, device=device)
    with config_context(array_api_dispatch=True):
        result_xp = _median(X_xp, axis=axis)

        if xp.__name__ != "array_api_strict":
            # We covert array-api-strict arrays to numpy arrays as `median` is not
            # part of the Array API spec
            assert get_namespace(result_xp)[0] == xp
            assert result_xp.device == X_xp.device
    assert_allclose(result_np, _convert_to_numpy(result_xp, xp=xp))


Review comment (suggestion):
```suggestion
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.
```

Code after:
    with config_context(array_api_dispatch=True):
        _fill_or_add_to_diagonal(array_xp, value=1, xp=xp, add_value=False, wrap=wrap)

    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])) is None
        assert get_namespace_and_device(a, b)[2] is None
        assert get_namespace_and_device(a, numpy.array([1]))[2] is 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)
    with config_context(array_api_dispatch=True):
        result_xp = _median(X_xp, axis=axis)

        if xp.__name__ != "array_api_strict":
            # We covert array-api-strict arrays to numpy arrays as `median` is not
            # part of the Array API spec
            assert get_namespace(result_xp)[0] == xp
            assert result_xp.device == X_xp.device
    assert_allclose(result_np, _convert_to_numpy(result_xp, xp=xp))