From f79de386b4d8fc45b50416bef41e1ac93dc2177e Mon Sep 17 00:00:00 2001 From: robinw Date: Tue, 28 Mar 2023 13:32:42 +0800 Subject: [PATCH 1/2] Fix concatenate and *stack APIs to support scalars --- cunumeric/module.py | 31 ++++++++-- tests/integration/test_concatenate_stack.py | 67 ++++++++++++++++++--- 2 files changed, 82 insertions(+), 16 deletions(-) diff --git a/cunumeric/module.py b/cunumeric/module.py index d6e84a8fa..ef0b27211 100644 --- a/cunumeric/module.py +++ b/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, diff --git a/tests/integration/test_concatenate_stack.py b/tests/integration/test_concatenate_stack.py index 33a11f392..52aba1672 100644 --- a/tests/integration/test_concatenate_stack.py +++ b/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 From 77b610af4d5f5586b3b7873504909688502d03fa Mon Sep 17 00:00:00 2001 From: robinw Date: Thu, 30 Mar 2023 14:30:21 +0800 Subject: [PATCH 2/2] Address comments --- cunumeric/module.py | 15 ++++++--------- 1 file changed, 6 insertions(+), 9 deletions(-) diff --git a/cunumeric/module.py b/cunumeric/module.py index ef0b27211..cfe0a7b4e 100644 --- a/cunumeric/module.py +++ b/cunumeric/module.py @@ -1286,7 +1286,7 @@ def _reshape_recur(ndim: int, arr: ndarray) -> tuple[int, ...]: def _atleast_nd( - ndim: int, arys: tuple[ndarray, ...] + 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 @@ -1806,7 +1806,7 @@ def concatenate( # flatten arrays if axis == None and concatenate arrays on the first axis if axis is None: # Reshape arrays in the `array_list` to handle scalars - reshaped = _atleast_nd(1, tuple(inputs)) + reshaped = _atleast_nd(1, inputs) if not isinstance(reshaped, list): reshaped = [reshaped] inputs = list(inp.ravel() for inp in reshaped) @@ -1875,9 +1875,6 @@ 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] @@ -1919,7 +1916,7 @@ def vstack(tup: Sequence[ndarray]) -> ndarray: Multiple GPUs, Multiple CPUs """ # Reshape arrays in the `array_list` if needed before concatenation - reshaped = _atleast_nd(2, tuple(tup)) + reshaped = _atleast_nd(2, tup) if not isinstance(reshaped, list): reshaped = [reshaped] tup, common_info = check_shape_dtype_without_axis( @@ -1968,7 +1965,7 @@ def hstack(tup: Sequence[ndarray]) -> ndarray: Multiple GPUs, Multiple CPUs """ # Reshape arrays in the `array_list` to handle scalars - reshaped = _atleast_nd(1, tuple(tup)) + reshaped = _atleast_nd(1, tup) if not isinstance(reshaped, list): reshaped = [reshaped] @@ -2022,7 +2019,7 @@ 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 = _atleast_nd(3, tuple(tup)) + reshaped = _atleast_nd(3, tup) if not isinstance(reshaped, list): reshaped = [reshaped] tup, common_info = check_shape_dtype_without_axis( @@ -2067,7 +2064,7 @@ 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)) + reshaped = _atleast_nd(1, tup) if not isinstance(reshaped, list): reshaped = [reshaped]