diff --git a/CHANGELOG.md b/CHANGELOG.md index 7ab890f77c..fc4cf6dc32 100644 --- a/CHANGELOG.md +++ b/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 diff --git a/dpctl/tensor/__init__.py b/dpctl/tensor/__init__.py index 1026774947..de5532c5ed 100644 --- a/dpctl/tensor/__init__.py +++ b/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", diff --git a/dpctl/tensor/_manipulation_functions.py b/dpctl/tensor/_manipulation_functions.py index 5beefab5ec..2e36b26dc1 100644 --- a/dpctl/tensor/_manipulation_functions.py +++ b/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 diff --git a/dpctl/tests/test_usm_ndarray_manipulation.py b/dpctl/tests/test_usm_ndarray_manipulation.py index 1b3716846b..9e99372639 100644 --- a/dpctl/tests/test_usm_ndarray_manipulation.py +++ b/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))