Move tests in test_common.py to test_helpers.py

Author: asmeurer

tests/test_common.py

@@ -5,6 +5,8 @@
 from ._helpers import import_
 
 import pytest
+import numpy as np
+from numpy.testing import assert_allclose
 
 is_functions = {
     'numpy': 'is_numpy_array',
@@ -44,3 +46,28 @@ def test_device(library):
 
     x2 = to_device(x, dev)
     assert device(x) == device(x2)
+
+
+@pytest.mark.parametrize("library", ["cupy", "numpy", "torch", "dask.array", "jax.numpy"])
+def test_to_device_host(library):
+    # Test that "cpu" device works. Note: this isn't actually supported by the
+    # standard yet. See https://github.com/data-apis/array-api/issues/626.
+
+    # different libraries have different semantics
+    # for DtoH transfers; ensure that we support a portable
+    # shim for common array libs
+    # see: https://github.com/scipy/scipy/issues/18286#issuecomment-1527552919
+    if library == "jax.numpy":
+        xp = import_('jax.experimental.array_api')
+    else:
+        xp = import_('array_api_compat.' + library)
+
+    expected = np.array([1, 2, 3])
+    x = xp.asarray([1, 2, 3])
+    x = to_device(x, "cpu")
+    # torch will return a genuine Device object, but
+    # the other libs will do something different with
+    # a `device(x)` query; however, what's really important
+    # here is that we can test portably after calling
+    # to_device(x, "cpu") to return to host
+    assert_allclose(x, expected)