Added support for argument axis with value None for dpctl.tensor.concat().

Author: npolina4

dpctl/tensor/_manipulation_functions.py

@@ -321,10 +321,10 @@ def roll(X, shift, axis=None):
     return res
 
 
-def _arrays_validation(arrays):
+def _arrays_validation(arrays, check_ndim=True):
     n = len(arrays)
     if n == 0:
-        raise TypeError("Missing 1 required positional argument: 'arrays'")
+        raise TypeError("Missing 1 required positional argument: 'arrays'.")
 
     if not isinstance(arrays, (list, tuple)):
         raise TypeError(f"Expected tuple or list type, got {type(arrays)}.")
@@ -335,11 +335,11 @@ def _arrays_validation(arrays):
 
     exec_q = dputils.get_execution_queue([X.sycl_queue for X in arrays])
     if exec_q is None:
-        raise ValueError("All the input arrays must have same sycl queue")
+        raise ValueError("All the input arrays must have same sycl queue.")
 
     res_usm_type = dputils.get_coerced_usm_type([X.usm_type for X in arrays])
     if res_usm_type is None:
-        raise ValueError("All the input arrays must have usm_type")
+        raise ValueError("All the input arrays must have usm_type.")
 
     X0 = arrays[0]
     _supported_dtype(Xi.dtype for Xi in arrays)
@@ -348,13 +348,14 @@ def _arrays_validation(arrays):
     for i in range(1, n):
         res_dtype = np.promote_types(res_dtype, arrays[i])
 
-    for i in range(1, n):
-        if X0.ndim != arrays[i].ndim:
-            raise ValueError(
-                "All the input arrays must have same number of dimensions, "
-                f"but the array at index 0 has {X0.ndim} dimension(s) and the "
-                f"array at index {i} has {arrays[i].ndim} dimension(s)"
-            )
+    if check_ndim:
+        for i in range(1, n):
+            if X0.ndim != arrays[i].ndim:
+                raise ValueError(
+                    "All the input arrays must have same number of dimensions, "
+                    f"but the array at index 0 has {X0.ndim} dimension(s) and "
+                    f"the array at index {i} has {arrays[i].ndim} dimension(s)."
+                )
     return res_dtype, res_usm_type, exec_q
 
 
@@ -367,7 +368,7 @@ def _check_same_shapes(X0_shape, axis, n, arrays):
                     "All the input array dimensions for the concatenation "
                     f"axis must match exactly, but along dimension {j}, the "
                     f"array at index 0 has size {X0j} and the array "
-                    f"at index {i} has size {Xi_shape[j]}"
+                    f"at index {i} has size {Xi_shape[j]}."
                 )
 
 
@@ -377,42 +378,66 @@ def concat(arrays, axis=0):
 
     Joins a sequence of arrays along an existing axis.
     """
-    res_dtype, res_usm_type, exec_q = _arrays_validation(arrays)
-
-    n = len(arrays)
-    X0 = arrays[0]
+    if axis is None:
+        res_dtype, res_usm_type, exec_q = _arrays_validation(
+            arrays, check_ndim=False
+        )
+        res_shape = 0
+        for array in arrays:
+            res_shape += array.size
+        res = dpt.empty(
+            res_shape, dtype=res_dtype, usm_type=res_usm_type, sycl_queue=exec_q
+        )
 
-    axis = normalize_axis_index(axis, X0.ndim)
-    X0_shape = X0.shape
-    _check_same_shapes(X0_shape, axis, n, arrays)
+        hev_list = []
+        fill_start = 0
+        for array in arrays:
+            fill_end = fill_start + array.size
+            hev, _ = ti._copy_usm_ndarray_into_usm_ndarray(
+                src=dpt.reshape(array, -1),
+                dst=res[fill_start:fill_end],
+                sycl_queue=exec_q,
+            )
+            fill_start = fill_end
+            hev_list.append(hev)
 
-    res_shape_axis = 0
-    for X in arrays:
-        res_shape_axis = res_shape_axis + X.shape[axis]
+        dpctl.SyclEvent.wait_for(hev_list)
+    else:
+        res_dtype, res_usm_type, exec_q = _arrays_validation(arrays)
+        n = len(arrays)
+        X0 = arrays[0]
 
-    res_shape = tuple(
-        X0_shape[i] if i != axis else res_shape_axis for i in range(X0.ndim)
-    )
+        axis = normalize_axis_index(axis, X0.ndim)
+        X0_shape = X0.shape
+        _check_same_shapes(X0_shape, axis, n, arrays)
 
-    res = dpt.empty(
-        res_shape, dtype=res_dtype, usm_type=res_usm_type, sycl_queue=exec_q
-    )
+        res_shape_axis = 0
+        for X in arrays:
+            res_shape_axis = res_shape_axis + X.shape[axis]
 
-    hev_list = []
-    fill_start = 0
-    for i in range(n):
-        fill_end = fill_start + arrays[i].shape[axis]
-        c_shapes_copy = tuple(
-            np.s_[fill_start:fill_end] if j == axis else np.s_[:]
-            for j in range(X0.ndim)
+        res_shape = tuple(
+            X0_shape[i] if i != axis else res_shape_axis for i in range(X0.ndim)
         )
-        hev, _ = ti._copy_usm_ndarray_into_usm_ndarray(
-            src=arrays[i], dst=res[c_shapes_copy], sycl_queue=exec_q
+
+        res = dpt.empty(
+            res_shape, dtype=res_dtype, usm_type=res_usm_type, sycl_queue=exec_q
         )
-        fill_start = fill_end
-        hev_list.append(hev)
 
-    dpctl.SyclEvent.wait_for(hev_list)
+        hev_list = []
+        fill_start = 0
+        for i in range(n):
+            fill_end = fill_start + arrays[i].shape[axis]
+            c_shapes_copy = tuple(
+                np.s_[fill_start:fill_end] if j == axis else np.s_[:]
+                for j in range(X0.ndim)
+            )
+            hev, _ = ti._copy_usm_ndarray_into_usm_ndarray(
+                src=arrays[i], dst=res[c_shapes_copy], sycl_queue=exec_q
+            )
+            fill_start = fill_end
+            hev_list.append(hev)
+
+        dpctl.SyclEvent.wait_for(hev_list)
 
     return res
 

dpctl/tests/test_usm_ndarray_manipulation.py

@@ -839,6 +839,7 @@ def test_concat_1array(data):
         [(0, 2), (2, 2), 0],
         [(2, 1), (2, 2), -1],
         [(2, 2, 2), (2, 1, 2), 1],
+        [(3, 3, 3), (2, 2), None],
     ],
 )
 def test_concat_2arrays(data):