Multi-device support for array manipulation

Author: betatim

array_api_strict/_manipulation_functions.py

@@ -25,8 +25,11 @@ def concat(
     # Note: Casting rules here are different from the np.concatenate default
     # (no for scalars with axis=None, no cross-kind casting)
     dtype = result_type(*arrays)
+    if len({a.device for a in arrays}) > 1:
+        raise ValueError("concat inputs must all be on the same device")
+
     arrays = tuple(a._array for a in arrays)
-    return Array._new(np.concatenate(arrays, axis=axis, dtype=dtype._np_dtype))
+    return Array._new(np.concatenate(arrays, axis=axis, dtype=dtype._np_dtype), device=arrays[0].device)
 
 
 def expand_dims(x: Array, /, *, axis: int) -> Array:
@@ -35,7 +38,7 @@ def expand_dims(x: Array, /, *, axis: int) -> Array:
 
     See its docstring for more information.
     """
-    return Array._new(np.expand_dims(x._array, axis))
+    return Array._new(np.expand_dims(x._array, axis), device=x.device)
 
 
 def flip(x: Array, /, *, axis: Optional[Union[int, Tuple[int, ...]]] = None) -> Array:
@@ -44,7 +47,7 @@ def flip(x: Array, /, *, axis: Optional[Union[int, Tuple[int, ...]]] = None) ->
 
     See its docstring for more information.
     """
-    return Array._new(np.flip(x._array, axis=axis))
+    return Array._new(np.flip(x._array, axis=axis), device=x.device)
 
 @requires_api_version('2023.12')
 def moveaxis(
@@ -58,7 +61,7 @@ def moveaxis(
 
     See its docstring for more information.
     """
-    return Array._new(np.moveaxis(x._array, source, destination))
+    return Array._new(np.moveaxis(x._array, source, destination), device=x.device)
 
 # Note: The function name is different here (see also matrix_transpose).
 # Unlike transpose(), the axes argument is required.
@@ -68,7 +71,7 @@ def permute_dims(x: Array, /, axes: Tuple[int, ...]) -> Array:
 
     See its docstring for more information.
     """
-    return Array._new(np.transpose(x._array, axes))
+    return Array._new(np.transpose(x._array, axes), device=x.device)
 
 @requires_api_version('2023.12')
 def repeat(
@@ -94,7 +97,7 @@ def repeat(
     else:
         raise TypeError("repeats must be an int or array")
 
-    return Array._new(np.repeat(x._array, repeats, axis=axis))
+    return Array._new(np.repeat(x._array, repeats, axis=axis), device=x.device)
 
 # Note: the optional argument is called 'shape', not 'newshape'
 def reshape(x: Array,
@@ -117,7 +120,7 @@ def reshape(x: Array,
     if copy is False and not np.shares_memory(data, reshaped):
         raise AttributeError("Incompatible shape for in-place modification.")
 
-    return Array._new(reshaped)
+    return Array._new(reshaped, device=x.device)
 
 
 def roll(
@@ -132,7 +135,7 @@ def roll(
 
     See its docstring for more information.
     """
-    return Array._new(np.roll(x._array, shift, axis=axis))
+    return Array._new(np.roll(x._array, shift, axis=axis), device=x.device)
 
 
 def squeeze(x: Array, /, axis: Union[int, Tuple[int, ...]]) -> Array:
@@ -141,7 +144,7 @@ def squeeze(x: Array, /, axis: Union[int, Tuple[int, ...]]) -> Array:
 
     See its docstring for more information.
     """
-    return Array._new(np.squeeze(x._array, axis=axis))
+    return Array._new(np.squeeze(x._array, axis=axis), device=x.device)
 
 
 def stack(arrays: Union[Tuple[Array, ...], List[Array]], /, *, axis: int = 0) -> Array:
@@ -152,8 +155,10 @@ def stack(arrays: Union[Tuple[Array, ...], List[Array]], /, *, axis: int = 0) ->
     """
     # Call result type here just to raise on disallowed type combinations
     result_type(*arrays)
+    if len({a.device for a in arrays}) > 1:
+        raise ValueError("concat inputs must all be on the same device")
     arrays = tuple(a._array for a in arrays)
-    return Array._new(np.stack(arrays, axis=axis))
+    return Array._new(np.stack(arrays, axis=axis), device=arrays[0].device)
 
 
 @requires_api_version('2023.12')
@@ -166,7 +171,7 @@ def tile(x: Array, repetitions: Tuple[int, ...], /) -> Array:
     # Note: NumPy allows repetitions to be an int or array
     if not isinstance(repetitions, tuple):
         raise TypeError("repetitions must be a tuple")
-    return Array._new(np.tile(x._array, repetitions))
+    return Array._new(np.tile(x._array, repetitions), device=x.device)
 
 # Note: this function is new
 @requires_api_version('2023.12')