From 57e0dce8bdae58c2b3322b77f84bcb49b9a7725b Mon Sep 17 00:00:00 2001 From: Aditya Goel Date: Thu, 29 Aug 2024 13:40:07 +0100 Subject: [PATCH 1/3] Fix cumsum with dtype specified --- ndonnx/_core/_numericimpl.py | 30 ++++++++++++++++-------------- ndonnx/_data_types/coretype.py | 6 +++++- tests/test_core.py | 26 +++++++++++++++++++++----- 3 files changed, 42 insertions(+), 20 deletions(-) diff --git a/ndonnx/_core/_numericimpl.py b/ndonnx/_core/_numericimpl.py index 885c920..1f83fdc 100644 --- a/ndonnx/_core/_numericimpl.py +++ b/ndonnx/_core/_numericimpl.py @@ -568,18 +568,17 @@ def cumulative_sum( else: raise ValueError("axis must be specified for multi-dimensional arrays") - if dtype is None: - if isinstance(x.dtype, (dtypes.Unsigned, dtypes.NullableUnsigned)): - if ndx.iinfo(x.dtype).bits < 64: - out = x.astype(dtypes.int64) - else: - raise ndx.UnsupportedOperationError( - f"Cannot perform `cumulative_sum` using {x.dtype}" - ) + if isinstance(x.dtype, (dtypes.Unsigned, dtypes.NullableUnsigned)): + if ndx.iinfo(x.dtype).bits < 64: + out = x.astype(dtypes.int64) else: - out = x.astype(_determine_reduce_op_dtype(x, dtype, dtypes.int64)) + return NotImplemented + elif dtype == dtypes.uint64 or dtype == dtypes.nuint64: + raise ndx.UnsupportedOperationError( + f"Unsupported dtype parameter for cumulative_sum {dtype} due to missing kernel support" + ) else: - out = out.astype(dtype) + out = x.astype(_determine_reduce_op_dtype(x, None, dtypes.int64)) out = from_corearray( opx.cumsum( @@ -589,10 +588,13 @@ def cumulative_sum( ) ) - if isinstance(x.dtype, dtypes.Unsigned): - out = out.astype(ndx.uint64) - elif isinstance(x.dtype, dtypes.NullableUnsigned): - out = out.astype(ndx.nuint64) + if dtype is None: + if isinstance(x.dtype, dtypes.Unsigned): + out = out.astype(ndx.uint64) + elif isinstance(x.dtype, dtypes.NullableUnsigned): + out = out.astype(ndx.nuint64) + else: + out = out.astype(dtype) # Exclude axis and create zeros of that shape if include_initial: diff --git a/ndonnx/_data_types/coretype.py b/ndonnx/_data_types/coretype.py index c82ee72..34fcdf5 100644 --- a/ndonnx/_data_types/coretype.py +++ b/ndonnx/_data_types/coretype.py @@ -69,7 +69,11 @@ def _parse_input(self, data: np.ndarray) -> dict[str, np.ndarray]: def _cast_from(self, array: Array) -> Array: if isinstance(array.dtype, CoreType): - return ndx.Array._from_fields(self, data=opx.cast(array._core(), to=self)) + return ( + ndx.Array._from_fields(self, data=opx.cast(array._core(), to=self)) + if array.dtype != self + else array.copy() + ) else: raise CastError(f"Cannot cast from {array.dtype} to {self}") diff --git a/tests/test_core.py b/tests/test_core.py index 589df5b..b0eaae3 100644 --- a/tests/test_core.py +++ b/tests/test_core.py @@ -933,13 +933,14 @@ def test_dynamic_reshape_has_no_static_shape(x, shape): ) @pytest.mark.parametrize("include_initial", [True, False]) @pytest.mark.parametrize( - "dtype", + "array_dtype", [ndx.int32, ndx.int64, ndx.float32, ndx.float64, ndx.uint8, ndx.uint16, ndx.uint32], ) @pytest.mark.parametrize( "array, axis", [ ([1, 2, 3], None), + ([100, 100], None), ([1, 2, 3], 0), ([[1, 2], [3, 4]], 0), ([[1, 2], [3, 4]], 1), @@ -948,21 +949,36 @@ def test_dynamic_reshape_has_no_static_shape(x, shape): ([[[[1]]], [[[3]]]], 1), ], ) -def test_cumulative_sum(array, axis, include_initial, dtype): - a = ndx.asarray(array, dtype=dtype) +@pytest.mark.parametrize( + "cumsum_dtype", + [None, ndx.int32, ndx.float32, ndx.float64, ndx.uint8, ndx.int8], +) +def test_cumulative_sum(array, axis, include_initial, array_dtype, cumsum_dtype): + a = ndx.asarray(array, dtype=array_dtype) assert_array_equal( - ndx.cumulative_sum(a, include_initial=include_initial, axis=axis).to_numpy(), + ndx.cumulative_sum( + a, include_initial=include_initial, axis=axis, dtype=cumsum_dtype + ).to_numpy(), np.cumulative_sum( np.asarray(array, a.dtype.to_numpy_dtype()), include_initial=include_initial, axis=axis, + dtype=cumsum_dtype.to_numpy_dtype() if cumsum_dtype is not None else None, ), ) def test_no_unsafe_cumulative_sum_cast(): with pytest.raises( - ndx.UnsupportedOperationError, match="Cannot perform `cumulative_sum`" + ndx.UnsupportedOperationError, + match="Unsupported operand type for cumulative_sum", ): a = ndx.asarray([1, 2, 3], ndx.uint64) ndx.cumulative_sum(a) + + with pytest.raises( + ndx.UnsupportedOperationError, + match="Unsupported dtype parameter for cumulative_sum", + ): + a = ndx.asarray([1, 2, 3], ndx.int32) + ndx.cumulative_sum(a, dtype=ndx.uint64) From ec451fdf1e50c11791a9c45d554aa9e7611d681b Mon Sep 17 00:00:00 2001 From: Aditya Goel Date: Thu, 29 Aug 2024 18:11:18 +0100 Subject: [PATCH 2/3] Use unsigned dtype --- ndonnx/_core/_numericimpl.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/ndonnx/_core/_numericimpl.py b/ndonnx/_core/_numericimpl.py index 1f83fdc..a3c92e7 100644 --- a/ndonnx/_core/_numericimpl.py +++ b/ndonnx/_core/_numericimpl.py @@ -578,7 +578,7 @@ def cumulative_sum( f"Unsupported dtype parameter for cumulative_sum {dtype} due to missing kernel support" ) else: - out = x.astype(_determine_reduce_op_dtype(x, None, dtypes.int64)) + out = x.astype(_determine_reduce_op_dtype(x, None, dtypes.uint64)) out = from_corearray( opx.cumsum( From e974def6bd73417323c56ed79a1d461b89e87476 Mon Sep 17 00:00:00 2001 From: Aditya Goel Date: Fri, 30 Aug 2024 09:33:02 +0100 Subject: [PATCH 3/3] Update change log date --- CHANGELOG.rst | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/CHANGELOG.rst b/CHANGELOG.rst index 00ffe0d..e12e1a2 100644 --- a/CHANGELOG.rst +++ b/CHANGELOG.rst @@ -8,7 +8,7 @@ Changelog ========= -0.9.0 (2024-08-29) +0.9.0 (2024-08-30) ------------------ **New features**