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

Code before:
        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)

        assert get_namespace(result_xp)[0].__name__ == xp.__name__
        assert result_xp.device == X_xp.device
        assert_allclose(result_np, _convert_to_numpy(result_xp, xp=xp))


Review comment (refactor):
💭 (this looks so ugly, I think because I don't love the look of `__foo__` 🙄, which made me wonder if we can't make it nicer. For no other reason than that it looks ugly, so take with many grains of salt. I had two thoughts: (1) make a helper like `matching_namespace(x, y)` that checks namespace and device (layers of indirection are tricky because you need to hold all of them in your head :-/) or (2) call `get_namespace(x, y)` which will raise an exception if `x` and `y` aren't in the same namespace (is it surprising to "abuse" `get_namespace` for this?))

Code after:
@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))