diff --git a/cunumeric/module.py b/cunumeric/module.py index 9da466f5d4..81a68309ab 100644 --- a/cunumeric/module.py +++ b/cunumeric/module.py @@ -1138,6 +1138,123 @@ def moveaxis( # Changing number of dimensions +def _reshape_recur(ndim: int, arr: ndarray) -> tuple[int]: + if arr.ndim < ndim: + cur_shape = _reshape_recur(ndim - 1, arr) + if ndim == 2: + cur_shape = (1,) + cur_shape + else: + cur_shape = cur_shape + (1,) + else: + cur_shape = arr.shape + return cur_shape + + +def _atleast_nd( + ndim: int, arys: Sequence[ndarray] +) -> Union(list[ndarray], ndarray): + inputs = list(convert_to_cunumeric_ndarray(arr) for arr in arys) + # 'reshape' change the shape of arrays + # only when arr.shape != _reshape_recur(ndim,arr) + result = list(arr.reshape(_reshape_recur(ndim, arr)) for arr in inputs) + # if the number of arrys in `arys` is 1, the return value is a single array + if len(result) == 1: + result = result[0] + return result + + +def atleast_1d(*arys: Sequence[ndarray]) -> Union(list[ndarray], ndarray): + """ + + Convert inputs to arrays with at least one dimension. + Scalar inputs are converted to 1-dimensional arrays, + whilst higher-dimensional inputs are preserved. + + Parameters + ---------- + *arys : array_like + One or more input arrays. + + Returns + ------- + ret : ndarray + An array, or list of arrays, each with a.ndim >= 1. + Copies are made only if necessary. + + See Also + -------- + numpy.atleast_1d + + Availability + -------- + Multiple GPUs, Multiple CPUs + """ + return _atleast_nd(1, arys) + + +def atleast_2d(*arys: Sequence[ndarray]) -> Union(list[ndarray], ndarray): + """ + + View inputs as arrays with at least two dimensions. + + Parameters + ---------- + *arys : array_like + One or more array-like sequences. + Non-array inputs are converted to arrays. + Arrays that already have two or more dimensions are preserved. + + Returns + ------- + res, res2, … : ndarray + An array, or list of arrays, each with a.ndim >= 2. + Copies are avoided where possible, and + views with two or more dimensions are returned. + + See Also + -------- + numpy.atleast_2d + + Availability + -------- + Multiple GPUs, Multiple CPUs + """ + return _atleast_nd(2, arys) + + +def atleast_3d(*arys: Sequence[ndarray]) -> Union(list[ndarray], ndarray): + """ + + View inputs as arrays with at least three dimensions. + + Parameters + ---------- + *arys : array_like + One or more array-like sequences. + Non-array inputs are converted to arrays. + Arrays that already have three or more dimensions are preserved. + + Returns + ------- + res, res2, … : ndarray + An array, or list of arrays, each with a.ndim >= 3. + Copies are avoided where possible, and + views with three or more dimensions are returned. + For example, a 1-D array of shape (N,) becomes + a view of shape (1, N, 1), and a 2-D array of shape (M, N) + becomes a view of shape (M, N, 1). + + See Also + -------- + numpy.atleast_3d + + Availability + -------- + Multiple GPUs, Multiple CPUs + """ + return _atleast_nd(3, arys) + + @add_boilerplate("a") def squeeze(a: ndarray, axis: Optional[NdShapeLike] = None) -> ndarray: """ @@ -1595,8 +1712,7 @@ def stack( " of input arrays" ) - shape = list(common_info.shape) - shape.insert(axis, 1) + shape = common_info.shape[:axis] + (1,) + common_info.shape[axis:] arrays = [arr.reshape(shape) for arr in arrays] common_info.shape = tuple(shape) return _concatenate(arrays, common_info, axis, out=out) @@ -1636,14 +1752,10 @@ def vstack(tup: Sequence[ndarray]) -> ndarray: Multiple GPUs, Multiple CPUs """ # Reshape arrays in the `array_list` if needed before concatenation - inputs = list(convert_to_cunumeric_ndarray(inp) for inp in tup) - reshaped = list( - inp.reshape([1, inp.shape[0]]) if inp.ndim == 1 else inp - for inp in inputs - ) + reshaped = _atleast_nd(2, tup) + if not isinstance(reshaped, list): + reshaped = [reshaped] tup, common_info = check_shape_dtype(reshaped, vstack.__name__, 0) - common_info.shape = tup[0].shape - return _concatenate( tup, common_info, @@ -1730,16 +1842,10 @@ def dstack(tup: Sequence[ndarray]) -> ndarray: Multiple GPUs, Multiple CPUs """ # Reshape arrays to (1,N,1) for ndim ==1 or (M,N,1) for ndim == 2: - reshaped = [] - inputs = list(convert_to_cunumeric_ndarray(inp) for inp in tup) - for arr in inputs: - if arr.ndim == 1: - arr = arr.reshape((1,) + arr.shape + (1,)) - elif arr.ndim == 2: - arr = arr.reshape(arr.shape + (1,)) - reshaped.append(arr) + reshaped = _atleast_nd(3, tup) + if not isinstance(reshaped, list): + reshaped = [reshaped] tup, common_info = check_shape_dtype(reshaped, dstack.__name__, 2) - return _concatenate( tup, common_info, diff --git a/docs/cunumeric/source/api/manipulation.rst b/docs/cunumeric/source/api/manipulation.rst index c930792050..bed086f10b 100644 --- a/docs/cunumeric/source/api/manipulation.rst +++ b/docs/cunumeric/source/api/manipulation.rst @@ -40,7 +40,9 @@ Changing number of dimensions :toctree: generated/ squeeze - + atleast_1d + atleast_2d + atleast_3d Changing kind of array ---------------------- diff --git a/tests/integration/test_atleast_nd.py b/tests/integration/test_atleast_nd.py new file mode 100644 index 0000000000..4cc011d637 --- /dev/null +++ b/tests/integration/test_atleast_nd.py @@ -0,0 +1,87 @@ +# Copyright 2022 NVIDIA Corporation +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. +# + +import numpy as np +import pytest + +import cunumeric as num +from legate.core import LEGATE_MAX_DIM + + +def _check(a, routine, sizes): + b = getattr(np, routine)(*a) + c = getattr(num, routine)(*a) + is_equal = True + err_arr = [b, c] + + if len(b) != len(c): + is_equal = False + err_arr = [b, c] + else: + for each in zip(b, c): + if not np.array_equal(*each): + err_arr = each + is_equal = False + break + print_msg = f"np.{routine}({sizes})" + assert is_equal, ( + f"Failed, {print_msg}\n" + f"numpy result: {err_arr[0]}\n" + f"cunumeric_result: {err_arr[1]}\n" + f"cunumeric and numpy shows different result\n" + ) + print(f"Passed, {print_msg}, np: {b}, cunumeric: {c}") + + +DIM = 10 + +SIZE_CASES = list((DIM,) * ndim for ndim in range(LEGATE_MAX_DIM + 1)) + +SIZE_CASES += [ + (0,), # empty array + (1,), # singlton array +] + + +# test to run atleast_nd w/ a single array +@pytest.mark.parametrize("size", SIZE_CASES, ids=str) +def test_atleast_1d(size): + a = [np.arange(np.prod(size)).reshape(size)] + _check(a, "atleast_1d", size) + + +@pytest.mark.parametrize("size", SIZE_CASES, ids=str) +def test_atleast_2d(size): + a = [np.arange(np.prod(size)).reshape(size)] + _check(a, "atleast_2d", size) + + +@pytest.mark.parametrize("size", SIZE_CASES, ids=str) +def test_atleast_3d(size): + a = [np.arange(np.prod(size)).reshape(size)] + _check(a, "atleast_3d", size) + + +# test to run atleast_nd w/ list of arrays +@pytest.mark.parametrize("dim", range(1, 4)) +def test_atleast_nd(dim): + a = list(np.arange(np.prod(size)).reshape(size) for size in SIZE_CASES) + _check(a, f"atleast_{dim}d", SIZE_CASES) + + +if __name__ == "__main__": + import sys + + sys.exit(pytest.main(sys.argv)) diff --git a/tests/integration/test_concatenate_stack.py b/tests/integration/test_concatenate_stack.py index 5d3e7be084..6ecc5b45ec 100644 --- a/tests/integration/test_concatenate_stack.py +++ b/tests/integration/test_concatenate_stack.py @@ -68,6 +68,8 @@ def run_test(arr, routine, input_size): DIM = 10 +NUM_ARR = [1, 3] + SIZES = [ (0,), (0, 10), @@ -81,40 +83,46 @@ def run_test(arr, routine, input_size): @pytest.fixture(autouse=True) -def a(size): - return [np.random.randint(low=0, high=100, size=size) for _ in range(3)] +def a(size, num): + return [np.random.randint(low=0, high=100, size=size) for _ in range(num)] +@pytest.mark.parametrize("num", NUM_ARR, ids=str) @pytest.mark.parametrize("size", SIZES, ids=str) -def test_concatenate(size, a): +def test_concatenate(size, num, a): run_test(tuple(a), "concatenate", size) +@pytest.mark.parametrize("num", NUM_ARR, ids=str) @pytest.mark.parametrize("size", SIZES, ids=str) -def test_stack(size, a): +def test_stack(size, num, a): run_test(tuple(a), "stack", size) +@pytest.mark.parametrize("num", NUM_ARR, ids=str) @pytest.mark.parametrize("size", SIZES, ids=str) -def test_hstack(size, a): +def test_hstack(size, num, a): run_test(tuple(a), "hstack", size) +@pytest.mark.parametrize("num", NUM_ARR, ids=str) @pytest.mark.parametrize("size", SIZES, ids=str) -def test_column_stack(size, a): +def test_column_stack(size, num, a): run_test(tuple(a), "column_stack", size) +@pytest.mark.parametrize("num", NUM_ARR, ids=str) @pytest.mark.parametrize("size", SIZES, ids=str) -def test_column_vstack(size, a): +def test_vstack(size, num, a): # exception for 1d array on vstack if len(size) == 2 and size == (1, DIM): a.append(np.random.randint(low=0, high=100, size=(DIM,))) run_test(tuple(a), "vstack", size) +@pytest.mark.parametrize("num", NUM_ARR, ids=str) @pytest.mark.parametrize("size", SIZES, ids=str) -def test_column_dstack(size, a): +def test_dstack(size, num, a): # exception for 1d array on dstack if len(size) == 2 and size == (1, DIM): a.append(np.random.randint(low=0, high=100, size=(DIM,)))