Merge pull request #867 from IntelPython/adding-concat-feature

Added dpctl.tensor.concat feature and tests

Author: oleksandr-pavlyk

CHANGELOG.md

@@ -31,6 +31,7 @@ and this project adheres to [Semantic Versioning](https://semver.org/spec/v2.0.0
 * Provided pybind11 example for functions working on `dpctl.tensor.usm_ndarray` container applying oneMKL functions [#780](https://github.com/IntelPython/dpctl/pull/780), [#793](https://github.com/IntelPython/dpctl/pull/793), [#819](https://github.com/IntelPython/dpctl/pull/819). The example was expanded to demonstrate implementing iterative linear solvers (Chebyshev solver, and Conjugate-Gradient solver) by asynchronously submitting individual SYCL kernels from Python [#821](https://github.com/IntelPython/dpctl/pull/821), [#833](https://github.com/IntelPython/dpctl/pull/833), [#838](https://github.com/IntelPython/dpctl/pull/838).
 * Wrote manual page about working with `dpctl.SyclQueue` [#829](https://github.com/IntelPython/dpctl/pull/829).
 * Added cmake scripts to dpctl package layout and a way to query the location [#853](https://github.com/IntelPython/dpctl/pull/853).
+* Implemented `dpctl.tensor.concat` function from array-API [#867](https://github.com/IntelPython/dpctl/867).
 
 
 ### Changed

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,90 @@ 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, 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 not all(Xi.dtype.char in "?bBhHiIlLqQefdFD" for Xi in arrays):
+        raise ValueError("Unsupported dtype encountered.")
+
+    res_dtype = X0.dtype
+    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, 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=res_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,167 @@ 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_different_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((3, 2), dtype=np.uint32, sycl_queue=q)
+
+    XY = dpt.concat([X, Y])
+
+    assert XY.dtype is X.dtype
+    assert XY.shape == (5, 2)
+    assert XY.sycl_queue == q
+
+
+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))