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

Code before:
    """Test get_namespace for ArrayAPI arrays."""
    xp = pytest.importorskip("array_api_strict")

    X_np = numpy.asarray([[1, 2, 3]])
    X_xp = xp.asarray(X_np)
    with config_context(array_api_dispatch=True):
        xp_out, is_array_api_compliant = get_namespace(X_xp)
        assert is_array_api_compliant

        with pytest.raises(TypeError):
            xp_out, is_array_api_compliant = get_namespace(X_xp, X_np)

        def mock_getenv(key):
            if key == "SCIPY_ARRAY_API":
                return "0"

        monkeypatch.setattr("os.environ.get", mock_getenv)
        assert os.environ.get("SCIPY_ARRAY_API") != "1"
        with pytest.raises(
            RuntimeError,
            match="scipy's own support is not enabled.",
        ):
            get_namespace(X_xp)


def test_move_to_array_api_conversions():
    """Check conversion of cupy and numpy to torch."""
    xp_torch = _array_api_for_tests("torch", "cuda")
    device_torch = xp_torch.asarray([1], device="cuda").device

    xp_cupy = _array_api_for_tests("cupy", None)
    array_cupy = xp_cupy.asarray([1, 2, 3], device=None)

    array_np = numpy.asarray([1, 2, 3], device=None)

    array_1_out, array_2_out = move_to(
        array_cupy, array_np, xp_reference=xp_torch, device_referenceerence=device_torch
    )
    for array in (array_1_out, array_2_out):
        assert get_namespace(array) == xp_torch
        assert device(array) == device_torch


def test_move_to_sparse():
    """Check sparse inputs are handled correctly."""
    xp_numpy = _array_api_for_tests("numpy", None)
    xp_torch = _array_api_for_tests("torch", "cpu")
    device_cpu = xp_torch.asarray([1]).device

    sparse1 = sp.csr_matrix([0, 1, 2, 3])
    sparse2 = sp.csr_matrix([0, 1, 0, 1])

Review comment (suggestion):
Maybe rename this to `test_move_to_array_api_conversions_with_cuda` and add a new similar function with torch + mps (that can run on our regular CI) and another with array-api-strict and its fake devices.

Code after:
@skip_if_array_api_compat_not_configured
def test_get_namespace_array_api(monkeypatch):
    """Test get_namespace for ArrayAPI arrays."""
    xp = pytest.importorskip("array_api_strict")

    X_np = numpy.asarray([[1, 2, 3]])
    X_xp = xp.asarray(X_np)
    with config_context(array_api_dispatch=True):
        xp_out, is_array_api_compliant = get_namespace(X_xp)
        assert is_array_api_compliant

        with pytest.raises(TypeError):
            xp_out, is_array_api_compliant = get_namespace(X_xp, X_np)

        def mock_getenv(key):
            if key == "SCIPY_ARRAY_API":
                return "0"

        monkeypatch.setattr("os.environ.get", mock_getenv)
        assert os.environ.get("SCIPY_ARRAY_API") != "1"
        with pytest.raises(
            RuntimeError,
            match="scipy's own support is not enabled.",
        ):
            get_namespace(X_xp)


@pytest.mark.parametrize(
    "array_input, reference",
    [
        pytest.param(("cupy", None), ("torch", "cuda"), id="cupy to torch cuda"),
        pytest.param(("torch", "mps"), ("numpy", None), id="torch mps to numpy"),
        pytest.param(("numpy", None), ("torch", "cuda"), id="numpy to torch cuda"),
        pytest.param(("numpy", None), ("torch", "mps"), id="numpy to torch mps"),
        pytest.param(
            ("array_api_strict", None),
            ("torch", "mps"),
            id="array_api_strict to torch mps",
        ),
    ],
)
def test_move_to_array_api_conversions(array_input, reference):
    """Check conversion between various namespace and devices."""
    if array_input[0] == "array_api_strict":
        array_api_strict = pytest.importorskip(
            "array_api_strict", reason="array-api-strict not available"
        )
    xp = _array_api_for_tests(reference[0], reference[1])
    xp_array = _array_api_for_tests(array_input[0], array_input[1])

    with config_context(array_api_dispatch=True):