Fix concatenate and *stack APIs to support scalars(#818, #839)

cunumeric/module.py

@@ -1479,8 +1479,8 @@ def check_shape_with_axis(
     ndim = inputs[0].ndim
     shape = inputs[0].shape
 
+    axis = normalize_axis_index(axis, ndim)
     if ndim >= 1:
-        axis = normalize_axis_index(axis, ndim)
         if _builtin_any(
             shape[:axis] != inp.shape[:axis]
             or shape[axis + 1 :] != inp.shape[axis + 1 :]
@@ -1805,7 +1805,11 @@ def concatenate(
 
     # flatten arrays if axis == None and concatenate arrays on the first axis
     if axis is None:
-        inputs = list(inp.ravel() for inp in inputs)
+        # Reshape arrays in the `array_list` to handle scalars
+        reshaped = _atleast_nd(1, tuple(inputs))
+        if not isinstance(reshaped, list):
+            reshaped = [reshaped]
+        inputs = list(inp.ravel() for inp in reshaped)
         axis = 0
 
     # Check to see if we can build a new tuple of cuNumeric arrays
@@ -1871,6 +1875,9 @@ def stack(
     if len(shapes) != 1:
         raise ValueError("all input arrays must have the same shape for stack")
 
+    # handle scalar inputs
+    if type(common_info.ndim) is not int:
+        common_info.ndim = 0
     axis = normalize_axis_index(axis, common_info.ndim + 1)
     shape = common_info.shape[:axis] + (1,) + common_info.shape[axis:]
     arrays = [arr.reshape(shape) for arr in arrays]
@@ -1960,7 +1967,14 @@ def hstack(tup: Sequence[ndarray]) -> ndarray:
     --------
     Multiple GPUs, Multiple CPUs
     """
-    tup, common_info = check_shape_dtype_without_axis(tup, hstack.__name__)
+    # Reshape arrays in the `array_list` to handle scalars
+    reshaped = _atleast_nd(1, tuple(tup))
+    if not isinstance(reshaped, list):
+        reshaped = [reshaped]
+
+    tup, common_info = check_shape_dtype_without_axis(
+        reshaped, hstack.__name__
+    )
     check_shape_with_axis(
         tup, hstack.__name__, axis=(0 if common_info.ndim == 1 else 1)
     )
@@ -2052,14 +2066,19 @@ def column_stack(tup: Sequence[ndarray]) -> ndarray:
     --------
     Multiple GPUs, Multiple CPUs
     """
+    # Reshape arrays in the `array_list` to handle scalars
+    reshaped = _atleast_nd(1, tuple(tup))
+    if not isinstance(reshaped, list):
+        reshaped = [reshaped]
+
     tup, common_info = check_shape_dtype_without_axis(
-        tup, column_stack.__name__
+        reshaped, column_stack.__name__
     )
-    # When ndim == 1, hstack concatenates arrays along the first axis
+
     if common_info.ndim == 1:
         tup = list(inp.reshape((inp.shape[0], 1)) for inp in tup)
         common_info.shape = tup[0].shape
-    check_shape_with_axis(tup, dstack.__name__, 1)
+    check_shape_with_axis(tup, column_stack.__name__, 1)
     return _concatenate(
         tup,
         common_info,

tests/integration/test_concatenate_stack.py

@@ -74,9 +74,7 @@ def run_test(arr, routine, input_size):
 NUM_ARR = [1, 3]
 
 SIZES = [
-    # In Numpy, hstack and column_stack PASS
-    # In cuNumeric, hstack and column_stack raise IndexError
-    pytest.param((), marks=pytest.mark.xfail),  # for scalar.
+    (),
     (0,),
     (0, 10),
     (1,),
@@ -87,6 +85,11 @@ def run_test(arr, routine, input_size):
     (DIM, DIM, DIM),
 ]
 
+SCALARS = (
+    (10,),
+    (10, 20, 30),
+)
+
 
 @pytest.fixture(autouse=False)
 def a(size, num):
@@ -99,6 +102,13 @@ def test_concatenate(size, num, a):
     run_test(tuple(a), "concatenate", size)
 
 
+@pytest.mark.parametrize("arrays", SCALARS, ids=str)
+def test_concatenate_scalar(arrays):
+    res_np = np.concatenate(arrays, axis=None)
+    res_num = num.concatenate(arrays, axis=None)
+    assert np.array_equal(res_np, res_num)
+
+
 def test_concatenate_with_out():
     a = [[1, 2], [3, 4]]
     b = [[5, 6]]
@@ -158,16 +168,13 @@ def test_zero_arrays(self):
     @pytest.mark.parametrize(
         "arrays",
         (
-            pytest.param((1,), marks=pytest.mark.xfail),
-            pytest.param((1, 2), marks=pytest.mark.xfail),
+            (1,),
+            (1, 2),
             (1, [3, 4]),
         ),
         ids=lambda arrays: f"(arrays={arrays})",
     )
     def test_scalar_axis_is_not_none(self, arrays):
-        # For (1,) and (1, 2),
-        # In Numpy, it raises ValueError
-        # In cuNumeric, it raises IndexError
         expected_exc = ValueError
         axis = 0
         with pytest.raises(expected_exc):
@@ -228,8 +235,6 @@ def test_both_out_dtype_are_provided(self):
             )
 
     def test_invalid_casting(self):
-        # In Numpy, raise ValueError
-        # In cuNumeric, pass
         expected_exc = ValueError
         a = [[1, 2], [3, 4]]
         b = [[5, 6]]
@@ -251,6 +256,13 @@ def test_stack(size, num, a):
     run_test(tuple(a), "stack", size)
 
 
+@pytest.mark.parametrize("arrays", SCALARS, ids=str)
+def test_stack_scalar(arrays):
+    res_np = np.stack(arrays)
+    res_num = num.stack(arrays)
+    assert np.array_equal(res_np, res_num)
+
+
 def test_stack_with_out():
     a = [1, 2]
     b = [3, 4]
@@ -351,6 +363,13 @@ def test_hstack(size, num, a):
     run_test(tuple(a), "hstack", size)
 
 
+@pytest.mark.parametrize("arrays", SCALARS, ids=str)
+def test_hstack_scalar(arrays):
+    res_np = np.hstack(arrays)
+    res_num = num.hstack(arrays)
+    assert np.array_equal(res_np, res_num)
+
+
 class TestHStackErrors:
     def test_zero_arrays(self):
         expected_exc = ValueError
@@ -382,6 +401,13 @@ def test_column_stack(size, num, a):
     run_test(tuple(a), "column_stack", size)
 
 
+@pytest.mark.parametrize("arrays", SCALARS, ids=str)
+def test_column_stack_scalar(arrays):
+    res_np = np.column_stack(arrays)
+    res_num = num.column_stack(arrays)
+    assert np.array_equal(res_np, res_num)
+
+
 class TestColumnStackErrors:
     def test_zero_arrays(self):
         expected_exc = ValueError
@@ -418,6 +444,13 @@ def test_vstack(size, num, a):
     run_test(tuple(a), "vstack", size)
 
 
+@pytest.mark.parametrize("arrays", SCALARS, ids=str)
+def test_vstack_scalar(arrays):
+    res_np = np.vstack(arrays)
+    res_num = num.vstack(arrays)
+    assert np.array_equal(res_np, res_num)
+
+
 class TestVStackErrors:
     def test_zero_arrays(self):
         expected_exc = ValueError
@@ -454,6 +487,13 @@ def test_rowstack(size, num, a):
     run_test(tuple(a), "row_stack", size)
 
 
+@pytest.mark.parametrize("arrays", SCALARS, ids=str)
+def test_row_stack_scalar(arrays):
+    res_np = np.row_stack(arrays)
+    res_num = num.row_stack(arrays)
+    assert np.array_equal(res_np, res_num)
+
+
 class TestRowStackErrors:
     def test_zero_arrays(self):
         expected_exc = ValueError
@@ -490,6 +530,13 @@ def test_dstack(size, num, a):
     run_test(tuple(a), "dstack", size)
 
 
+@pytest.mark.parametrize("arrays", SCALARS, ids=str)
+def test_dstack_scalar(arrays):
+    res_np = np.dstack(arrays)
+    res_num = num.dstack(arrays)
+    assert np.array_equal(res_np, res_num)
+
+
 class TestDStackErrors:
     def test_zero_arrays(self):
         expected_exc = ValueError