Adding dpctl.tensor.concat feature and tests.

concat() function concatenates several arrays along one of axis.
https://data-apis.org/array-api/latest/API_specification/generated/signatures.manipulation_functions.concat.html

Author: npolina4

dpctl/tensor/__init__.py

@@ -39,6 +39,7 @@
 from dpctl.tensor._manipulation_functions import (
     broadcast_arrays,
     broadcast_to,
+    concat,
     expand_dims,
     flip,
     permute_dims,
@@ -66,6 +67,7 @@
     "flip",
     "reshape",
     "roll",
+    "concat",
     "broadcast_arrays",
     "broadcast_to",
     "expand_dims",

dpctl/tensor/_manipulation_functions.py

@@ -18,11 +18,12 @@
 from itertools import chain, product, repeat
 
 import numpy as np
-from numpy.core.numeric import normalize_axis_tuple
+from numpy.core.numeric import normalize_axis_index, normalize_axis_tuple
 
 import dpctl
 import dpctl.tensor as dpt
 import dpctl.tensor._tensor_impl as ti
+import dpctl.utils as dputils
 
 
 def _broadcast_strides(X_shape, X_strides, res_ndim):
@@ -285,3 +286,86 @@ def roll(X, shift, axes=None):
 
     dpctl.SyclEvent.wait_for(hev_list)
     return res
+
+
+def concat(arrays, axis=0):
+    """
+    concat(arrays: tuple or list of usm_ndarrays, axis: int) -> usm_ndarray
+
+    Joins a sequence of arrays along an existing axis.
+    """
+    n = len(arrays)
+    if n == 0:
+        raise TypeError("Missing 1 required positional argument: 'arrays'")
+
+    if not isinstance(arrays, list) and not isinstance(arrays, tuple):
+        raise TypeError(f"Expected tuple or list type, got {type(arrays)}.")
+
+    for X in arrays:
+        if not isinstance(X, dpt.usm_ndarray):
+            raise TypeError(f"Expected usm_ndarray type, got {type(X)}.")
+
+    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")
+
+    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")
+
+    X0 = arrays[0]
+    if any(X0.dtype != arrays[i].dtype for i in range(1, n)):
+        raise ValueError("All the input arrays must have same dtype")
+
+    for i in range(1, n):
+        if X0.ndim != arrays[i].ndim:
+            raise ValueError(
+                "All the input arrays must have same number of "
+                "dimensions, but the array at index 0 has "
+                f"{X0.ndim} dimension(s) and the array at index "
+                f"{i} has {arrays[i].ndim} dimension(s)"
+            )
+
+    axis = normalize_axis_index(axis, X0.ndim)
+    X0_shape = X0.shape
+    for i in range(1, n):
+        Xi_shape = arrays[i].shape
+        for j in range(X0.ndim):
+            if X0_shape[j] != Xi_shape[j] and j != axis:
+                raise ValueError(
+                    "All the input array dimensions for the "
+                    "concatenation axis must match exactly, but "
+                    f"along dimension {j}, the array at index 0 "
+                    f"has size {X0_shape[j]} and the array at "
+                    f"index {i} has size {Xi_shape[j]}"
+                )
+
+    res_shape_axis = 0
+    for X in arrays:
+        res_shape_axis = res_shape_axis + X.shape[axis]
+
+    res_shape = tuple(
+        X0_shape[i] if i != axis else res_shape_axis for i in range(X0.ndim)
+    )
+
+    res = dpt.empty(
+        res_shape, dtype=X0.dtype, usm_type=res_usm_type, sycl_queue=exec_q
+    )
+
+    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

@@ -725,3 +725,163 @@ def test_roll_2d(data):
     Y = dpt.roll(X, sh, ax)
     Ynp = np.roll(Xnp, sh, ax)
     assert_array_equal(Ynp, dpt.asnumpy(Y))
+
+
+def test_concat_incorrect_type():
+    Xnp = np.ones((2, 2))
+    pytest.raises(TypeError, dpt.concat)
+    pytest.raises(TypeError, dpt.concat, [])
+    pytest.raises(TypeError, dpt.concat, Xnp)
+    pytest.raises(TypeError, dpt.concat, [Xnp, Xnp])
+
+
+def test_concat_incorrect_queue():
+    try:
+        q1 = dpctl.SyclQueue()
+        q2 = dpctl.SyclQueue()
+    except dpctl.SyclQueueCreationError:
+        pytest.skip("Queue could not be created")
+
+    X = dpt.ones((2, 2), sycl_queue=q1)
+    Y = dpt.ones((2, 2), sycl_queue=q2)
+
+    pytest.raises(ValueError, dpt.concat, [X, Y])
+
+
+def test_concat_incorrect_dtype():
+    try:
+        q = dpctl.SyclQueue()
+    except dpctl.SyclQueueCreationError:
+        pytest.skip("Queue could not be created")
+
+    X = dpt.ones((2, 2), dtype=np.int64, sycl_queue=q)
+    Y = dpt.ones((2, 2), dtype=np.uint64, sycl_queue=q)
+
+    pytest.raises(ValueError, dpt.concat, [X, Y])
+
+
+def test_concat_incorrect_ndim():
+    try:
+        q = dpctl.SyclQueue()
+    except dpctl.SyclQueueCreationError:
+        pytest.skip("Queue could not be created")
+
+    X = dpt.ones((2, 2), sycl_queue=q)
+    Y = dpt.ones((2, 2, 2), sycl_queue=q)
+
+    pytest.raises(ValueError, dpt.concat, [X, Y])
+
+
+@pytest.mark.parametrize(
+    "data",
+    [
+        [(2, 2), (3, 3), 0],
+        [(2, 2), (3, 3), 1],
+        [(3, 2), (3, 3), 0],
+        [(2, 3), (3, 3), 1],
+    ],
+)
+def test_concat_incorrect_shape(data):
+    try:
+        q = dpctl.SyclQueue()
+    except dpctl.SyclQueueCreationError:
+        pytest.skip("Queue could not be created")
+
+    Xshape, Yshape, axis = data
+
+    X = dpt.ones(Xshape, sycl_queue=q)
+    Y = dpt.ones(Yshape, sycl_queue=q)
+
+    pytest.raises(ValueError, dpt.concat, [X, Y], axis)
+
+
+@pytest.mark.parametrize(
+    "data",
+    [
+        [(6,), 0],
+        [(2, 3), 1],
+        [(3, 2), -1],
+        [(1, 6), 0],
+        [(2, 1, 3), 2],
+    ],
+)
+def test_concat_1array(data):
+    try:
+        q = dpctl.SyclQueue()
+    except dpctl.SyclQueueCreationError:
+        pytest.skip("Queue could not be created")
+
+    Xshape, axis = data
+
+    Xnp = np.arange(6).reshape(Xshape)
+    X = dpt.asarray(Xnp, sycl_queue=q)
+
+    Ynp = np.concatenate([Xnp], axis=axis)
+    Y = dpt.concat([X], axis=axis)
+
+    assert_array_equal(Ynp, dpt.asnumpy(Y))
+
+    Ynp = np.concatenate((Xnp,), axis=axis)
+    Y = dpt.concat((X,), axis=axis)
+
+    assert_array_equal(Ynp, dpt.asnumpy(Y))
+
+
+@pytest.mark.parametrize(
+    "data",
+    [
+        [(1,), (1,), 0],
+        [(0, 2), (2, 2), 0],
+        [(2, 1), (2, 2), -1],
+        [(2, 2, 2), (2, 1, 2), 1],
+    ],
+)
+def test_concat_2arrays(data):
+    try:
+        q = dpctl.SyclQueue()
+    except dpctl.SyclQueueCreationError:
+        pytest.skip("Queue could not be created")
+
+    Xshape, Yshape, axis = data
+
+    Xnp = np.ones(Xshape)
+    X = dpt.asarray(Xnp, sycl_queue=q)
+
+    Ynp = np.zeros(Yshape)
+    Y = dpt.asarray(Ynp, sycl_queue=q)
+
+    Znp = np.concatenate([Xnp, Ynp], axis=axis)
+    Z = dpt.concat([X, Y], axis=axis)
+
+    assert_array_equal(Znp, dpt.asnumpy(Z))
+
+
+@pytest.mark.parametrize(
+    "data",
+    [
+        [(1,), (1,), (1,), 0],
+        [(0, 2), (2, 2), (1, 2), 0],
+        [(2, 1, 2), (2, 2, 2), (2, 4, 2), 1],
+    ],
+)
+def test_concat_3arrays(data):
+    try:
+        q = dpctl.SyclQueue()
+    except dpctl.SyclQueueCreationError:
+        pytest.skip("Queue could not be created")
+
+    Xshape, Yshape, Zshape, axis = data
+
+    Xnp = np.ones(Xshape)
+    X = dpt.asarray(Xnp, sycl_queue=q)
+
+    Ynp = np.zeros(Yshape)
+    Y = dpt.asarray(Ynp, sycl_queue=q)
+
+    Znp = np.full(Zshape, 2.0)
+    Z = dpt.asarray(Znp, sycl_queue=q)
+
+    Rnp = np.concatenate([Xnp, Ynp, Znp], axis=axis)
+    R = dpt.concat([X, Y, Z], axis=axis)
+
+    assert_array_equal(Rnp, dpt.asnumpy(R))