BUG: fix cumsum/cumprod with arrays on a non-default device

Author: ev-br

array_api_strict/_statistical_functions.py

@@ -39,7 +39,10 @@ def cumulative_sum(
     if include_initial:
         if axis < 0:
             axis += x.ndim
-        x = concat([zeros(x.shape[:axis] + (1,) + x.shape[axis + 1:], dtype=x.dtype), x], axis=axis)
+        x = concat(
+            [zeros(x.shape[:axis] + (1,) + x.shape[axis + 1:], dtype=x.dtype, device=x.device), x],
+            axis=axis
+        )
     return Array._new(np.cumsum(x._array, axis=axis, dtype=_np_dtype(dtype)), device=x.device)
 
 
@@ -66,7 +69,10 @@ def cumulative_prod(
     if include_initial:
         if axis < 0:
             axis += x.ndim
-        x = concat([ones(x.shape[:axis] + (1,) + x.shape[axis + 1:], dtype=x.dtype), x], axis=axis)
+        x = concat(
+            [ones(x.shape[:axis] + (1,) + x.shape[axis + 1:], dtype=x.dtype, device=x.device), x],
+            axis=axis
+        )
     return Array._new(np.cumprod(x._array, axis=axis, dtype=_np_dtype(dtype)), device=x.device)
 
 

array_api_strict/tests/test_statistical_functions.py

@@ -55,3 +55,19 @@ def test_mean_complex():
     with pytest.raises(TypeError):
         xp.mean(xp.arange(3))
 
+
+def test_cumsum_device():
+    x = xp.arange(3, device=xp.Device('device1'))
+    y = xp.cumulative_sum(x, include_initial=True)
+    expected = xp.asarray([0, 0, 1, 3], device=x.device)
+    assert y.device == expected.device
+    assert xp.all(y == expected)
+
+
+def test_cumprod_device():
+    x = xp.arange(1, 4, device=xp.Device('device1'))
+    y = xp.cumulative_prod(x, include_initial=True)
+    expected = xp.asarray([1, 1, 2, 6], device=x.device)
+    assert y.device == expected.device
+    assert xp.all(y == expected)
+