diff --git a/doc/array.rst b/doc/array.rst index de1cc66ed3b1ecc09b4b999a95b04461215a9b55..34c65693e8c4244c5a6c401eb35d4573a56b0971 100644 --- a/doc/array.rst +++ b/doc/array.rst @@ -339,6 +339,15 @@ Constructing :class:`GPUArray` Instances Return the :class:`GPUArray` ``[a[indices[0]], ..., a[indices[n]]]``. For the moment, *a* must be a type that can be bound to a texture. +.. function:: concatenate(arrays, axis=0, allocator=None) + + Join a sequence of arrays along an existing axis. + +.. function:: stack(arrays, axis=0, allocator=None) + + Join a sequence of arrays along a new axis. + + Conditionals ^^^^^^^^^^^^ diff --git a/pycuda/gpuarray.py b/pycuda/gpuarray.py index 373cf00532b51c7e870d2e0c09e092e418e74378..a0bf84c4d42f93ba2215c895f7e3746a2d9b177a 100644 --- a/pycuda/gpuarray.py +++ b/pycuda/gpuarray.py @@ -1767,6 +1767,95 @@ def multi_put(arrays, dest_indices, dest_shape=None, out=None, stream=None): # {{{ shape manipulation +def concatenate(arrays, axis=0, allocator=None): + """ + Join a sequence of arrays along an existing axis. + :arg arrays: A sequnce of :class:`GPUArray`. + :arg axis: Index of the dimension of the new axis in the result array. + Can be -1, for the new axis to be last dimension. + :returns: :class:`GPUArray` + """ + # implementation is borrowed from pyopencl.array.concatenate() + # {{{ find properties of result array + + shape = None + + def shape_except_axis(ary: GPUArray): + return ary.shape[:axis] + ary.shape[axis+1:] + + for i_ary, ary in enumerate(arrays): + allocator = allocator or ary.allocator + + if shape is None: + # first array + shape = list(ary.shape) + + else: + if len(ary.shape) != len(shape): + raise ValueError("%d'th array has different number of axes " + "(should have %d, has %d)" + % (i_ary, len(ary.shape), len(shape))) + + if (ary.ndim != arrays[0].ndim + or shape_except_axis(ary) != shape_except_axis(arrays[0])): + raise ValueError("%d'th array has residual not matching " + "other arrays" % i_ary) + + shape[axis] += ary.shape[axis] + + # }}} + + shape = tuple(shape) + dtype = np.find_common_type([ary.dtype for ary in arrays], []) + result = empty(shape, dtype, allocator=allocator) + + full_slice = (slice(None),) * len(shape) + + base_idx = 0 + for ary in arrays: + my_len = ary.shape[axis] + result[full_slice[:axis] + (slice(base_idx, base_idx+my_len),) + full_slice[axis+1:]] = ary + base_idx += my_len + + return result + + +def stack(arrays, axis=0, allocator=None): + """ + Join a sequence of arrays along a new axis. + :arg arrays: A sequnce of :class:`GPUArray`. + :arg axis: Index of the dimension of the new axis in the result array. + Can be -1, for the new axis to be last dimension. + :returns: :class:`GPUArray` + """ + # implementation is borrowed from pyopencl.array.stack() + allocator = allocator or arrays[0].allocator + + if not arrays: + raise ValueError("need at least one array to stack") + + input_shape = arrays[0].shape + input_ndim = arrays[0].ndim + axis = input_ndim if axis == -1 else axis + + if not all(ary.shape == input_shape for ary in arrays[1:]): + raise ValueError("arrays must have the same shape") + + if not (0 <= axis <= input_ndim): + raise ValueError("invalid axis") + + result_shape = input_shape[:axis] + (len(arrays),) + input_shape[axis:] + result = empty(shape=result_shape, + dtype=np.result_type(*(ary.dtype for ary in arrays)), + allocator=allocator, order="C" if axis == 0 else "F") + + for i, ary in enumerate(arrays): + + idx = (slice(None),)*axis + (i,) + (slice(None),)*(input_ndim-axis) + result[idx] = ary + + return result + def transpose(a, axes=None): """Permute the dimensions of an array. diff --git a/test/test_gpuarray.py b/test/test_gpuarray.py index 4b53aca7ba62d8508eff6cf8829342529339a158..dbf4b2f7f2ebee4fbec4a810b2ff3f8437d99458 100644 --- a/test/test_gpuarray.py +++ b/test/test_gpuarray.py @@ -468,6 +468,53 @@ class TestGPUArray: a = gpuarray.arange(12, dtype=np.float32) assert (np.arange(12, dtype=np.float32) == a.get()).all() + @mark_cuda_test + def test_stack(self): + + orders = ["F", "C"] + input_dims_lst = [0, 1, 2] + + for order in orders: + for input_dims in input_dims_lst: + shape = (2, 2, 2)[:input_dims] + axis = -1 if order == "F" else 0 + + from numpy.random import default_rng + rng = default_rng() + x_in = rng.random(size=shape) + y_in = rng.random(size=shape) + x_in = x_in if order == "C" else np.asfortranarray(x_in) + y_in = y_in if order == "C" else np.asfortranarray(y_in) + + x_gpu = gpuarray.to_gpu(x_in) + y_gpu = gpuarray.to_gpu(y_in) + + numpy_stack = np.stack((x_in, y_in), axis=axis) + gpuarray_stack = gpuarray.stack((x_gpu, y_gpu), axis=axis) + + np.testing.assert_allclose(gpuarray_stack.get(), numpy_stack) + + assert gpuarray_stack.shape == numpy_stack.shape + + @mark_cuda_test + def test_concatenate(self): + + from pycuda.curandom import rand as curand + + a_dev = curand((5, 15, 20), dtype=np.float32) + b_dev = curand((4, 15, 20), dtype=np.float32) + c_dev = curand((3, 15, 20), dtype=np.float32) + a = a_dev.get() + b = b_dev.get() + c = c_dev.get() + + cat_dev = gpuarray.concatenate((a_dev, b_dev, c_dev)) + cat = np.concatenate((a, b, c)) + + np.testing.assert_allclose(cat, cat_dev.get()) + + assert cat.shape == cat_dev.shape + @mark_cuda_test def test_reverse(self): a = np.array([1, 2, 3, 4, 5, 6, 7, 8, 9, 10]).astype(np.float32)