Add a basic test for device() and to_device()

Author: asmeurer

tests/test_helpers.py

@@ -1,5 +1,6 @@
 from array_api_compat import (is_numpy_array, is_cupy_array, is_torch_array,
-                              is_dask_array, is_jax_array, is_array_api_obj)
+                              is_dask_array, is_jax_array, is_array_api_obj,
+                              device, to_device)
 
 from ._helpers import import_
 
@@ -24,3 +25,19 @@ def test_is_xp_array(library, func):
     assert is_func(x) == (func == is_functions[library])
 
     assert is_array_api_obj(x)
+
+@pytest.mark.parametrize("library", ["cupy", "numpy", "torch", "dask.array", "jax.numpy"])
+def test_device(library):
+    if library == "jax.numpy":
+        xp = import_('jax.experimental.array_api')
+    else:
+        xp = import_('array_api_compat.' + library)
+
+    # We can't test much for device() and to_device() other than that
+    # x.to_device(x.device) works.
+
+    x = xp.asarray([1, 2, 3])
+    dev = device(x)
+
+    x2 = to_device(x, dev)
+    assert device(x) == device(x2)