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

Code before:
    """Check conversion of Numpy to torch mps."""
    xp_torch = _array_api_for_tests("torch", "mps")
    device_torch = xp_torch.asarray([1], device="mps").device

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

    (array_out,) = move_to(
        array_np, xp_reference=xp_torch, device_reference=device_torch
    )

    with config_context(array_api_dispatch=True):
        assert get_namespace(array_out)[0] == xp_torch
        assert device(array_out) == device_torch


def test_move_to_array_api_conversions_with_strict():
    """Check conversion of array-api-strict, with different devices, to torch."""
    try:
        import array_api_strict

        array_api_strict_available = True
    except ImportError:
        array_api_strict_available = False

    if not array_api_strict_available:
        pytest.mark.skip(reason="array-api-strict not available")

    xp_strict = _array_api_for_tests("array_api_strict", None)
    array_strict_cpu = xp_strict.asarray(
        [1, 2, 3], device=array_api_strict.Device("CPU_DEVICE")
    )
    # array_strict_device1 = xp_strict.asarray(
    #     [1, 2, 3], device=array_api_strict.Device("device1")
    # )

    xp_torch = _array_api_for_tests("torch", "mps")
    device_torch = xp_torch.asarray([1], device="mps").device
    (array_1_out,) = move_to(
        array_strict_cpu,
        # array_strict_device1,
        xp_reference=xp_torch,
        device_reference=device_torch,
    )

    with config_context(array_api_dispatch=True):
        # for array in (array_1_out,):
        assert get_namespace(array_1_out)[0] == xp_torch
        assert device(array_1_out) == device_torch


def test_move_to_sparse():

Review comment (suggestion):
We could alternatively use the `@pytest.mark.skipif` decorator, i believe,  which is a bit shorter.

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

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

        # sparse and None to NumPy
        result1, result2 = move_to(sparse1, None, xp=xp_numpy, device=None)
        assert result1 is sparse1
        assert result2 is None

        # sparse to non-NumPy
        msg = r"Sparse arrays are only accepted \(and passed through\)"
        with pytest.raises(TypeError, match=msg):
            move_to(sparse1, numpy_array, xp=xp_torch, device=device_cpu)
        with pytest.raises(TypeError, match=msg):
            move_to(sparse1, None, xp=xp_torch, device=device_cpu)


@pytest.mark.parametrize("array_api", ["numpy", "array_api_strict"])
def test_asarray_with_order(array_api):
    """Test _asarray_with_order passes along order for NumPy arrays."""
    xp = pytest.importorskip(array_api)

    X = xp.asarray([1.2, 3.4, 5.1])
    X_new = _asarray_with_order(X, order="F", xp=xp)

    X_new_np = numpy.asarray(X_new)