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

Code before:
            match="scipy's own support is not enabled.",
        ):
            get_namespace(X_xp)


@pytest.mark.parametrize(
    "array_input, reference",
    (
        (("cupy", None), ("torch", "cuda")),
        (("torch", "mps"), ("numpy", None)),
        (("numpy", None), ("torch", "cuda")),
        (("numpy", None), ("torch", "mps")),
        (("array_api_strict", None), ("torch", "mps")),
    ),
)
def test_move_to_array_api_conversions(array_input, reference):
    """Check conversion between various namespace and devices."""
    if array_input == "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):
        device_ = xp.asarray([1], device=reference[1]).device

        if array_input == "array_api_strict":
            array_device = array_api_strict.Device("CPU_DEVICE")
        else:
            array_device = array_input[1]
        array = xp_array.asarray([1, 2, 3], device=array_device)

        array_out = move_to(array, xp=xp, device=device_)
        assert get_namespace(array_out)[0] == xp
        assert device(array_out) == device_


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")

    sparse1 = sp.csr_array([0, 1, 2, 3])
    sparse2 = sp.csr_array([0, 1, 0, 1])
    numpy_array = numpy.array([1, 2, 3])

    with config_context(array_api_dispatch=True):
        device_cpu = xp_torch.asarray([1]).device

        # sparse to NumPy

Review comment (suggestion):
```suggestion
        device_ = device(xp.asarray([1], device=reference[1]))
```

Maybe use scikit-learn's `device` function instead of calling the `.device` attribute?

```bash
import numpy as np

from sklearn import config_context
from sklearn.utils._array_api import device

with config_context(array_api_dispatch=True):
    y = [1, 2, 3, 3, 2, 1, 1, 2, 3]
    y_np = np.asarray(y)

    print(f"{device(y_np) = }") # None
    print(f"{y_np.device = }") # "cpu"
```

Code after:
            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):
        device_ = device(xp.asarray([1], device=reference[1]))

        if array_input[0] == "array_api_strict":
            array_device = array_api_strict.Device("CPU_DEVICE")
        else:
            array_device = array_input[1]
        array = xp_array.asarray([1, 2, 3], device=array_device)

        array_out = move_to(array, xp=xp, device=device_)
        assert get_namespace(array_out)[0] == xp
        assert device(array_out) == device_


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")

    sparse1 = sp.csr_array([0, 1, 2, 3])
    sparse2 = sp.csr_array([0, 1, 0, 1])
    numpy_array = numpy.array([1, 2, 3])