Implementation of atleast_{1,2,3}d

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,

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
 ----------------------

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))

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,)))