diff --git a/src/awkward/_do.py b/src/awkward/_do.py index a502ce1656..1458fba125 100644 --- a/src/awkward/_do.py +++ b/src/awkward/_do.py @@ -217,31 +217,34 @@ def pad_none( return layout._pad_none(length, axis, 1, clip) -def completely_flatten( +def remove_structure( layout: Content | Record, backend: Backend | None = None, flatten_records: bool = True, function_name: str | None = None, drop_nones: bool = True, + keepdims: bool = False, ): if isinstance(layout, Record): - return completely_flatten( + return remove_structure( layout._array[layout._at : layout._at + 1], backend, flatten_records, function_name, drop_nones, + keepdims, ) else: if backend is None: backend = layout._backend - arrays = layout._completely_flatten( + arrays = layout._remove_structure( backend, { "flatten_records": flatten_records, "function_name": function_name, "drop_nones": drop_nones, + "keepdims": keepdims, }, ) return tuple(arrays) @@ -314,15 +317,16 @@ def reduce( behavior: dict | None = None, ): if axis is None: - parts = completely_flatten(layout, flatten_records=False, drop_nones=False) + parts = remove_structure( + layout, flatten_records=False, drop_nones=False, keepdims=keepdims + ) if len(parts) > 1: # We know that `flatten_records` must fail, so the only other type # that can return multiple parts here is the union array raise ak._errors.wrap_error( ValueError( - "cannot use axis=None with keepdims=True on an array containing " - "irreducible unions" + "cannot use axis=None on an array containing irreducible unions" ) ) elif len(parts) == 0: diff --git a/src/awkward/_util.py b/src/awkward/_util.py index 5230e55304..7765449232 100644 --- a/src/awkward/_util.py +++ b/src/awkward/_util.py @@ -771,6 +771,7 @@ def arrays_approx_equal( atol: float = 1e-8, dtype_exact: bool = True, check_parameters=True, + check_regular=True, ) -> bool: # TODO: this should not be needed after refactoring nplike mechanism import awkward.forms.form @@ -798,7 +799,13 @@ def visitor(left, right) -> bool: right = right.to_IndexedOptionArray64() if type(left) is not type(right): - return False + if not check_regular and ( + left.is_list and right.is_regular or left.is_regular and right.is_list + ): + left = left.to_ListOffsetArray64() + right = right.to_ListOffsetArray64() + else: + return False if left.length != right.length: return False diff --git a/src/awkward/contents/bitmaskedarray.py b/src/awkward/contents/bitmaskedarray.py index 8e3f64bbd6..89b694fab7 100644 --- a/src/awkward/contents/bitmaskedarray.py +++ b/src/awkward/contents/bitmaskedarray.py @@ -585,10 +585,10 @@ def _to_arrow(self, pyarrow, mask_node, validbytes, length, options): def _to_backend_array(self, allow_missing, backend): return self.to_ByteMaskedArray()._to_backend_array(allow_missing, backend) - def _completely_flatten(self, backend, options): + def _remove_structure(self, backend, options): branch, depth = self.branch_depth if branch or options["drop_nones"] or depth > 1: - return self.project()._completely_flatten(backend, options) + return self.project()._remove_structure(backend, options) else: return [self] diff --git a/src/awkward/contents/bytemaskedarray.py b/src/awkward/contents/bytemaskedarray.py index 381926acfd..40468a8224 100644 --- a/src/awkward/contents/bytemaskedarray.py +++ b/src/awkward/contents/bytemaskedarray.py @@ -972,10 +972,10 @@ def _to_arrow(self, pyarrow, mask_node, validbytes, length, options): def _to_backend_array(self, allow_missing, backend): return self.to_IndexedOptionArray64()._to_backend_array(allow_missing, backend) - def _completely_flatten(self, backend, options): + def _remove_structure(self, backend, options): branch, depth = self.branch_depth if branch or options["drop_nones"] or depth > 1: - return self.project()._completely_flatten(backend, options) + return self.project()._remove_structure(backend, options) else: return [self] diff --git a/src/awkward/contents/content.py b/src/awkward/contents/content.py index c47f31cec1..10439dab47 100644 --- a/src/awkward/contents/content.py +++ b/src/awkward/contents/content.py @@ -1095,7 +1095,7 @@ def drop_none(self): def _drop_none(self) -> Content: raise ak._errors.wrap_error(NotImplementedError) - def _completely_flatten(self, backend, options): + def _remove_structure(self, backend, options): raise ak._errors.wrap_error(NotImplementedError) def _recursively_apply( diff --git a/src/awkward/contents/emptyarray.py b/src/awkward/contents/emptyarray.py index 7518e3be0a..a41323789b 100644 --- a/src/awkward/contents/emptyarray.py +++ b/src/awkward/contents/emptyarray.py @@ -308,7 +308,7 @@ def _to_arrow(self, pyarrow, mask_node, validbytes, length, options): def _to_backend_array(self, allow_missing, backend): return backend.nplike.empty(0, dtype=np.float64) - def _completely_flatten(self, backend, options): + def _remove_structure(self, backend, options): return [] def _recursively_apply( diff --git a/src/awkward/contents/indexedarray.py b/src/awkward/contents/indexedarray.py index e94ec61762..32eedcde53 100644 --- a/src/awkward/contents/indexedarray.py +++ b/src/awkward/contents/indexedarray.py @@ -958,8 +958,8 @@ def _to_arrow(self, pyarrow, mask_node, validbytes, length, options): def _to_backend_array(self, allow_missing, backend): return self.project()._to_backend_array(allow_missing, backend) - def _completely_flatten(self, backend, options): - return self.project()._completely_flatten(backend, options) + def _remove_structure(self, backend, options): + return self.project()._remove_structure(backend, options) def _recursively_apply( self, action, behavior, depth, depth_context, lateral_context, options diff --git a/src/awkward/contents/indexedoptionarray.py b/src/awkward/contents/indexedoptionarray.py index 20ab049661..f604f87626 100644 --- a/src/awkward/contents/indexedoptionarray.py +++ b/src/awkward/contents/indexedoptionarray.py @@ -1331,7 +1331,7 @@ def _reduce_next( "reduce_next with unbranching depth > negaxis is only " "expected to return RegularArray or ListOffsetArray or " "IndexedOptionArray; " - "instead, it returned " + out + "instead, it returned {}".format(type(out).__name__) ) ) @@ -1526,10 +1526,10 @@ def _to_backend_array(self, allow_missing, backend): else: return content - def _completely_flatten(self, backend, options): + def _remove_structure(self, backend, options): branch, depth = self.branch_depth if branch or options["drop_nones"] or depth > 1: - return self.project()._completely_flatten(backend, options) + return self.project()._remove_structure(backend, options) else: return [self] diff --git a/src/awkward/contents/listarray.py b/src/awkward/contents/listarray.py index 3841fc9e31..c78e53868c 100644 --- a/src/awkward/contents/listarray.py +++ b/src/awkward/contents/listarray.py @@ -1379,16 +1379,8 @@ def _to_arrow(self, pyarrow, mask_node, validbytes, length, options): def _to_backend_array(self, allow_missing, backend): return self.to_RegularArray()._to_backend_array(allow_missing, backend) - def _completely_flatten(self, backend, options): - if ( - self.parameter("__array__") == "string" - or self.parameter("__array__") == "bytestring" - ): - return [self] - else: - next = self.to_ListOffsetArray64(False) - flat = next.content[next.offsets[0] : next.offsets[-1]] - return flat._completely_flatten(backend, options) + def _remove_structure(self, backend, options): + return self.to_ListOffsetArray64(False)._remove_structure(backend, options) def _drop_none(self): return self.to_ListOffsetArray64()._drop_none() diff --git a/src/awkward/contents/listoffsetarray.py b/src/awkward/contents/listoffsetarray.py index c8048d8049..24781b8ed6 100644 --- a/src/awkward/contents/listoffsetarray.py +++ b/src/awkward/contents/listoffsetarray.py @@ -1962,15 +1962,30 @@ def _to_backend_array(self, allow_missing, backend): return self.to_RegularArray()._to_backend_array(allow_missing, backend) - def _completely_flatten(self, backend, options): + def _remove_structure(self, backend, options): if ( self.parameter("__array__") == "string" or self.parameter("__array__") == "bytestring" ): return [self] else: - flat = self._content[self._offsets[0] : self._offsets[-1]] - return flat._completely_flatten(backend, options) + content = self._content[self._offsets[0] : self._offsets[-1]] + contents = content._remove_structure(backend, options) + if options["keepdims"]: + return [ + ListOffsetArray( + ak.index.Index64( + backend.index_nplike.asarray( + [0, backend.index_nplike.shape_item_as_scalar(c.length)] + ) + ), + c, + parameters=self._parameters, + ) + for c in contents + ] + else: + return contents def _drop_none(self): if self._content.is_option: diff --git a/src/awkward/contents/numpyarray.py b/src/awkward/contents/numpyarray.py index e171f13cba..959b2b4e39 100644 --- a/src/awkward/contents/numpyarray.py +++ b/src/awkward/contents/numpyarray.py @@ -1204,10 +1204,14 @@ def _to_arrow(self, pyarrow, mask_node, validbytes, length, options): def _to_backend_array(self, allow_missing, backend): return to_nplike(self.data, backend.nplike, from_nplike=self._backend.nplike) - def _completely_flatten(self, backend, options): + def _remove_structure(self, backend, options): + if options["keepdims"]: + shape = (1,) * (self._data.ndim - 1) + (-1,) + else: + shape = (-1,) return [ ak.contents.NumpyArray( - backend.nplike.reshape(self._raw(backend.nplike), (-1,)), + backend.nplike.reshape(self._raw(backend.nplike), shape), backend=backend, ) ] diff --git a/src/awkward/contents/recordarray.py b/src/awkward/contents/recordarray.py index ff12a37757..91a078585c 100644 --- a/src/awkward/contents/recordarray.py +++ b/src/awkward/contents/recordarray.py @@ -922,13 +922,11 @@ def _to_backend_array(self, allow_missing, backend): return out - def _completely_flatten(self, backend, options): + def _remove_structure(self, backend, options): if options["flatten_records"]: out = [] for content in self._contents: - out.extend( - content[: self._length]._completely_flatten(backend, options) - ) + out.extend(content[: self._length]._remove_structure(backend, options)) return out else: in_function = "" diff --git a/src/awkward/contents/regulararray.py b/src/awkward/contents/regulararray.py index c810b6d7d9..545a3debc4 100644 --- a/src/awkward/contents/regulararray.py +++ b/src/awkward/contents/regulararray.py @@ -1218,7 +1218,7 @@ def _to_arrow(self, pyarrow, mask_node, validbytes, length, options): ), ) - def _completely_flatten(self, backend, options): + def _remove_structure(self, backend, options): if ( self.parameter("__array__") == "string" or self.parameter("__array__") == "bytestring" @@ -1227,8 +1227,15 @@ def _completely_flatten(self, backend, options): else: index_nplike = self._backend.index_nplike length = index_nplike.mul_shape_item(self._length, self._size) - flat = self._content[: index_nplike.shape_item_as_scalar(length)] - return flat._completely_flatten(backend, options) + content = self._content[: index_nplike.shape_item_as_scalar(length)] + contents = content._remove_structure(backend, options) + if options["keepdims"]: + return [ + RegularArray(c, size=c.length, parameters=self._parameters) + for c in contents + ] + else: + return contents def _drop_none(self): return self.to_ListOffsetArray64()._drop_none() diff --git a/src/awkward/contents/unionarray.py b/src/awkward/contents/unionarray.py index 8349df4728..c175c14e70 100644 --- a/src/awkward/contents/unionarray.py +++ b/src/awkward/contents/unionarray.py @@ -1469,14 +1469,14 @@ def _to_backend_array(self, allow_missing, backend): return out - def _completely_flatten(self, backend, options): + def _remove_structure(self, backend, options): out = [] for i in range(len(self._contents)): index = self._index[self._tags.data == i] out.extend( self._contents[i] ._carry(index, False) - ._completely_flatten(backend, options) + ._remove_structure(backend, options) ) return out diff --git a/src/awkward/contents/unmaskedarray.py b/src/awkward/contents/unmaskedarray.py index 301abd1ad7..b17cb3382f 100644 --- a/src/awkward/contents/unmaskedarray.py +++ b/src/awkward/contents/unmaskedarray.py @@ -424,10 +424,10 @@ def _to_backend_array(self, allow_missing, backend): else: return content - def _completely_flatten(self, backend, options): + def _remove_structure(self, backend, options): branch, depth = self.branch_depth if branch or options["drop_nones"] or depth > 1: - return self.project()._completely_flatten(backend, options) + return self.project()._remove_structure(backend, options) else: return [self] diff --git a/src/awkward/operations/ak_flatten.py b/src/awkward/operations/ak_flatten.py index 2bc018553f..089d1e0af1 100644 --- a/src/awkward/operations/ak_flatten.py +++ b/src/awkward/operations/ak_flatten.py @@ -166,7 +166,7 @@ def _impl(array, axis, highlevel, behavior): layout = ak.operations.to_layout(array, allow_record=False, allow_other=False) if axis is None: - out = ak._do.completely_flatten(layout, function_name="ak.flatten") + out = ak._do.remove_structure(layout, function_name="ak.flatten") assert isinstance(out, tuple) and all( isinstance(x, ak.contents.NumpyArray) for x in out ) diff --git a/src/awkward/operations/ak_ravel.py b/src/awkward/operations/ak_ravel.py index 7b5ee77442..484b4c6570 100644 --- a/src/awkward/operations/ak_ravel.py +++ b/src/awkward/operations/ak_ravel.py @@ -58,7 +58,7 @@ def ravel(array, *, highlevel=True, behavior=None): def _impl(array, highlevel, behavior): layout = ak.operations.to_layout(array, allow_record=False, allow_other=False) - out = ak._do.completely_flatten(layout, function_name="ak.ravel", drop_nones=False) + out = ak._do.remove_structure(layout, function_name="ak.ravel", drop_nones=False) assert isinstance(out, tuple) and all( isinstance(x, ak.contents.Content) for x in out ) diff --git a/tests/test_2020_reduce_axis_none.py b/tests/test_2020_reduce_axis_none.py index 9a07491e1a..803408131c 100644 --- a/tests/test_2020_reduce_axis_none.py +++ b/tests/test_2020_reduce_axis_none.py @@ -11,11 +11,11 @@ def test_sum(): assert ak.sum(array, axis=None) == pytest.approx(63.0) assert ak._util.arrays_approx_equal( - ak.sum(array, axis=None, keepdims=True), ak.Array([63.0]) + ak.sum(array, axis=None, keepdims=True), ak.to_regular([[63.0]]) ) assert ak._util.arrays_approx_equal( ak.sum(array, axis=None, keepdims=True, mask_identity=True), - ak.Array([63.0]).mask[[True]], + ak.to_regular(ak.Array([[63.0]]).mask[[[True]]]), ) assert ak.sum(array[2], axis=None, mask_identity=True) is None @@ -24,14 +24,14 @@ def test_prod(): assert ak.prod(array[1:], axis=None) == pytest.approx(4838400.0) assert ak.prod(array, axis=None) == 0 assert ak._util.arrays_approx_equal( - ak.prod(array, axis=None, keepdims=True), ak.Array([0.0]) + ak.prod(array, axis=None, keepdims=True), ak.to_regular([[0.0]]) ) assert ak._util.arrays_approx_equal( - ak.prod(array[1:], axis=None, keepdims=True), ak.Array([4838400.0]) + ak.prod(array[1:], axis=None, keepdims=True), ak.to_regular([[4838400.0]]) ) assert ak._util.arrays_approx_equal( ak.prod(array[1:], axis=None, keepdims=True, mask_identity=True), - ak.Array([4838400.0]).mask[[True]], + ak.to_regular(ak.Array([[4838400.0]]).mask[[[True]]]), ) assert ak.prod(array[2], axis=None, mask_identity=True) is None @@ -40,20 +40,20 @@ def test_min(): assert ak.min(array, axis=None) == pytest.approx(0.0) assert ak._util.arrays_approx_equal( ak.min(array, axis=None, keepdims=True, mask_identity=False), - ak.Array([0.0]), + ak.to_regular([[0.0]]), ) assert ak._util.arrays_approx_equal( ak.min(array, axis=None, keepdims=True, initial=-100, mask_identity=False), - ak.Array([-100.0]), + ak.to_regular([[-100.0]]), ) assert ak._util.arrays_approx_equal( ak.min(array, axis=None, keepdims=True, mask_identity=True), - ak.Array([0.0]).mask[[True]], + ak.to_regular(ak.Array([[0.0]]).mask[[[True]]]), ) assert ak._util.arrays_approx_equal( ak.min(array[-1:], axis=None, keepdims=True, mask_identity=True), - ak.Array([np.inf]).mask[[False]], + ak.to_regular(ak.Array([[np.inf]]).mask[[[False]]]), ) assert ak.min(array[2], axis=None, mask_identity=True) is None @@ -62,19 +62,19 @@ def test_max(): assert ak.max(array, axis=None) == pytest.approx(10.0) assert ak._util.arrays_approx_equal( ak.max(array, axis=None, keepdims=True, mask_identity=False), - ak.Array([10.0]), + ak.to_regular([[10.0]]), ) assert ak._util.arrays_approx_equal( ak.max(array, axis=None, keepdims=True, initial=100, mask_identity=False), - ak.Array([100.0]), + ak.to_regular([[100.0]]), ) assert ak._util.arrays_approx_equal( ak.max(array, axis=None, keepdims=True, mask_identity=True), - ak.Array([10.0]).mask[[True]], + ak.to_regular(ak.Array([[10.0]]).mask[[[True]]]), ) assert ak._util.arrays_approx_equal( ak.max(array[-1:], axis=None, keepdims=True, mask_identity=True), - ak.Array([np.inf]).mask[[False]], + ak.to_regular(ak.Array([[np.inf]]).mask[[[False]]]), ) assert ak.max(array[2], axis=None, mask_identity=True) is None @@ -83,15 +83,15 @@ def test_count(): assert ak.count(array, axis=None) == 12 assert ak._util.arrays_approx_equal( ak.count(array, axis=None, keepdims=True, mask_identity=False), - ak.Array([12]), + ak.to_regular([[12]]), ) assert ak._util.arrays_approx_equal( ak.count(array, axis=None, keepdims=True, mask_identity=True), - ak.Array([12]).mask[[True]], + ak.to_regular(ak.Array([[12]]).mask[[[True]]]), ) assert ak._util.arrays_approx_equal( ak.count(array[-1:], axis=None, keepdims=True, mask_identity=True), - ak.Array([0]).mask[[False]], + ak.to_regular(ak.Array([[0]]).mask[[[False]]]), ) assert ak.count(array[2], axis=None, mask_identity=True) is None assert ak.count(array[2], axis=None, mask_identity=False) == 0 @@ -101,15 +101,15 @@ def test_count_nonzero(): assert ak.count_nonzero(array, axis=None) == 11 assert ak._util.arrays_approx_equal( ak.count_nonzero(array, axis=None, keepdims=True, mask_identity=False), - ak.Array([11]), + ak.to_regular([[11]]), ) assert ak._util.arrays_approx_equal( ak.count_nonzero(array, axis=None, keepdims=True, mask_identity=True), - ak.Array([11]).mask[[True]], + ak.to_regular(ak.Array([[11]]).mask[[[True]]]), ) assert ak._util.arrays_approx_equal( ak.count_nonzero(array[-1:], axis=None, keepdims=True, mask_identity=True), - ak.Array([0]).mask[[False]], + ak.to_regular(ak.Array([[0]]).mask[[[False]]]), ) assert ak.count_nonzero(array[2], axis=None, mask_identity=True) is None assert ak.count_nonzero(array[2], axis=None, mask_identity=False) == 0 @@ -119,11 +119,11 @@ def test_std(): assert ak.std(array, axis=None) == pytest.approx(3.139134700306227) assert ak._util.arrays_approx_equal( ak.std(array, axis=None, keepdims=True, mask_identity=False), - ak.Array([3.139134700306227]), + ak.to_regular([[3.139134700306227]]), ) assert ak._util.arrays_approx_equal( ak.std(array, axis=None, keepdims=True, mask_identity=True), - ak.Array([3.139134700306227]).mask[[True]], + ak.to_regular(ak.Array([[3.139134700306227]]).mask[[[True]]]), ) assert np.isnan(ak.std(array[2], axis=None, mask_identity=False)) @@ -131,7 +131,7 @@ def test_std(): def test_std_no_mask_axis_none(): assert ak._util.arrays_approx_equal( ak.std(array[-1:], axis=None, keepdims=True, mask_identity=True), - ak.Array([0.0]).mask[[False]], + ak.to_regular(ak.Array([[0.0]]).mask[[[False]]]), ) assert ak.std(array[2], axis=None, mask_identity=True) is None @@ -140,11 +140,11 @@ def test_var(): assert ak.var(array, axis=None) == pytest.approx(9.854166666666666) assert ak._util.arrays_approx_equal( ak.var(array, axis=None, keepdims=True, mask_identity=False), - ak.Array([9.854166666666666]), + ak.to_regular([[9.854166666666666]]), ) assert ak._util.arrays_approx_equal( ak.var(array, axis=None, keepdims=True, mask_identity=True), - ak.Array([9.854166666666666]).mask[[True]], + ak.to_regular(ak.Array([[9.854166666666666]]).mask[[[True]]]), ) assert np.isnan(ak.var(array[2], axis=None, mask_identity=False)) @@ -152,7 +152,7 @@ def test_var(): def test_var_no_mask_axis_none(): assert ak._util.arrays_approx_equal( ak.var(array[-1:], axis=None, keepdims=True, mask_identity=True), - ak.Array([0.0]).mask[[False]], + ak.to_regular(ak.Array([[0.0]]).mask[[[False]]]), ) assert ak.var(array[2], axis=None, mask_identity=True) is None @@ -161,11 +161,11 @@ def test_mean(): assert ak.mean(array, axis=None) == pytest.approx(5.25) assert ak._util.arrays_approx_equal( ak.mean(array, axis=None, keepdims=True, mask_identity=False), - ak.Array([5.25]), + ak.to_regular([[5.25]]), ) assert ak._util.arrays_approx_equal( ak.mean(array, axis=None, keepdims=True, mask_identity=True), - ak.Array([5.25]).mask[[True]], + ak.to_regular(ak.Array([[5.25]]).mask[[[True]]]), ) assert np.isnan(ak.mean(array[2], axis=None, mask_identity=False)) @@ -173,7 +173,7 @@ def test_mean(): def test_mean_no_mask_axis_none(): assert ak._util.arrays_approx_equal( ak.mean(array[-1:], axis=None, keepdims=True, mask_identity=True), - ak.Array([0.0]).mask[[False]], + ak.to_regular(ak.Array([[0.0]]).mask[[[False]]]), ) assert ak.mean(array[2], axis=None, mask_identity=True) is None @@ -182,11 +182,11 @@ def test_ptp(): assert ak.ptp(array, axis=None) == pytest.approx(10.0) assert ak._util.arrays_approx_equal( ak.ptp(array, axis=None, keepdims=True, mask_identity=False), - ak.Array([10.0]), + ak.to_regular([[10.0]]), ) assert ak._util.arrays_approx_equal( ak.ptp(array, axis=None, keepdims=True, mask_identity=True), - ak.Array([10.0]).mask[[True]], + ak.to_regular(ak.Array([[10.0]]).mask[[[True]]]), ) assert ak.ptp(array[2], axis=None, mask_identity=False) == pytest.approx(0.0) @@ -194,7 +194,7 @@ def test_ptp(): def test_ptp_no_mask_axis_none(): assert ak._util.arrays_approx_equal( ak.ptp(array[-1:], axis=None, keepdims=True, mask_identity=True), - ak.Array([0.0]).mask[[False]], + ak.to_regular(ak.Array([[0.0]]).mask[[[False]]]), ) assert ak.ptp(array[2], axis=None, mask_identity=True) is None @@ -203,15 +203,15 @@ def test_argmax(): assert ak.argmax(array, axis=None) == 11 assert ak._util.arrays_approx_equal( ak.argmax(array, axis=None, keepdims=True, mask_identity=False), - ak.Array([11]), + ak.to_regular([[11]]), ) assert ak._util.arrays_approx_equal( ak.argmax(array, axis=None, keepdims=True, mask_identity=True), - ak.Array([11]).mask[[True]], + ak.to_regular(ak.Array([[11]]).mask[[[True]]]), ) assert ak._util.arrays_approx_equal( ak.argmax(array[-1:], axis=None, keepdims=True, mask_identity=True), - ak.Array([0]).mask[[False]], + ak.to_regular(ak.Array([[0]]).mask[[[False]]]), ) assert ak.argmax(array[2], axis=None, mask_identity=True) is None assert ak.argmax(array[2], axis=None, mask_identity=False) == -1 @@ -221,15 +221,15 @@ def test_argmin(): assert ak.argmin(array, axis=None) == 0 assert ak._util.arrays_approx_equal( ak.argmin(array, axis=None, keepdims=True, mask_identity=False), - ak.Array([0]), + ak.to_regular([[0]]), ) assert ak._util.arrays_approx_equal( ak.argmin(array, axis=None, keepdims=True, mask_identity=True), - ak.Array([0]).mask[[True]], + ak.to_regular(ak.Array([[0]]).mask[[[True]]]), ) assert ak._util.arrays_approx_equal( ak.argmin(array[-1:], axis=None, keepdims=True, mask_identity=True), - ak.Array([999]).mask[[False]], + ak.to_regular(ak.Array([[999]]).mask[[[False]]]), ) assert ak.argmin(array[2], axis=None, mask_identity=True) is None assert ak.argmin(array[2], axis=None, mask_identity=False) == -1 @@ -239,15 +239,15 @@ def test_any(): assert ak.any(array, axis=None) assert ak._util.arrays_approx_equal( ak.any(array, axis=None, keepdims=True, mask_identity=False), - ak.Array([True]), + ak.to_regular([[True]]), ) assert ak._util.arrays_approx_equal( ak.any(array, axis=None, keepdims=True, mask_identity=True), - ak.Array([True]).mask[[True]], + ak.to_regular(ak.Array([[True]]).mask[[[True]]]), ) assert ak._util.arrays_approx_equal( ak.any(array[-1:], axis=None, keepdims=True, mask_identity=True), - ak.Array([True]).mask[[False]], + ak.to_regular(ak.Array([[True]]).mask[[[False]]]), ) assert ak.any(array[2], axis=None, mask_identity=True) is None assert not ak.any(array[2], axis=None, mask_identity=False) @@ -257,15 +257,29 @@ def test_all(): assert not ak.all(array, axis=None) assert ak._util.arrays_approx_equal( ak.all(array, axis=None, keepdims=True, mask_identity=False), - ak.Array([False]), + ak.to_regular([[False]]), ) assert ak._util.arrays_approx_equal( ak.all(array, axis=None, keepdims=True, mask_identity=True), - ak.Array([False]).mask[[True]], + ak.to_regular(ak.Array([[False]]).mask[[[True]]]), ) assert ak._util.arrays_approx_equal( ak.all(array[-1:], axis=None, keepdims=True, mask_identity=True), - ak.Array([False]).mask[[False]], + ak.to_regular(ak.Array([[False]]).mask[[[False]]]), ) assert ak.all(array[2], axis=None, mask_identity=True) is None assert ak.all(array[2], axis=None, mask_identity=False) + + +def test_prod_numpy(): + array_reg = ak.from_numpy(np.arange(7 * 5 * 3, dtype=np.int64).reshape((7, 5, 3))) + result_reg = ak.from_numpy(np.array([[[5460]]], dtype=np.int64)) + + assert ak.sum(array_reg, axis=None) == 5460 + assert ak._util.arrays_approx_equal( + ak.sum(array_reg, axis=None, keepdims=True), result_reg + ) + assert ak._util.arrays_approx_equal( + ak.sum(array_reg, axis=None, keepdims=True, mask_identity=True), + ak.to_regular(result_reg.mask[[[[True]]]], axis=None), + )