add xp, device tests

Author: lucascolley

src/array_api_extra/_funcs.py

@@ -582,7 +582,10 @@ def pad(
         xp = array_namespace(x)
 
     padded = xp.full(
-        tuple(x + 2 * pad_width for x in x.shape), fill_value=value, dtype=x.dtype
+        tuple(x + 2 * pad_width for x in x.shape),
+        fill_value=value,
+        dtype=x.dtype,
+        device=_compat.device(x),
     )
     padded[(slice(pad_width, -pad_width, None),) * x.ndim] = x
     return padded

tests/test_funcs.py

@@ -408,3 +408,11 @@ def test_mode_not_implemented(self):
         a = xp.arange(3)
         with pytest.raises(NotImplementedError, match="Only `'constant'`"):
             pad(a, 2, mode="edge")
+
+    def test_device(self):
+        device = xp.Device("device1")
+        a = xp.asarray(0.0, device=device)
+        assert pad(a, 2).device == device
+
+    def test_xp(self):
+        assert_array_equal(pad(xp.asarray(0), 1, xp=xp), xp.zeros(3))