diff --git a/src/awkward/operations/ak_argcombinations.py b/src/awkward/operations/ak_argcombinations.py index 337f77cec1..0009bb1ad7 100644 --- a/src/awkward/operations/ak_argcombinations.py +++ b/src/awkward/operations/ak_argcombinations.py @@ -108,20 +108,17 @@ def _impl( axis = regularize_axis(axis, none_allowed=False) - if axis < 0: - raise ValueError("the 'axis' for argcombinations must be non-negative") - else: - with HighLevelContext(behavior=behavior, attrs=attrs) as ctx: - layout = ak._do.local_index( - ctx.unwrap(array, allow_record=False, primitive_policy="error"), - axis, - ) - out = ak._do.combinations( - layout, - n, - replacement=replacement, - axis=axis, - fields=fields, - parameters=parameters, + with HighLevelContext(behavior=behavior, attrs=attrs) as ctx: + layout = ak._do.local_index( + ctx.unwrap(array, allow_record=False, primitive_policy="error"), + axis, ) - return ctx.wrap(out, highlevel=highlevel) + out = ak._do.combinations( + layout, + n, + replacement=replacement, + axis=axis, + fields=fields, + parameters=parameters, + ) + return ctx.wrap(out, highlevel=highlevel) diff --git a/tests/test_3300_allow_negative_axis_in_argcombinations.py b/tests/test_3300_allow_negative_axis_in_argcombinations.py new file mode 100644 index 0000000000..1907a04293 --- /dev/null +++ b/tests/test_3300_allow_negative_axis_in_argcombinations.py @@ -0,0 +1,22 @@ +from __future__ import annotations + +import awkward as ak + + +def test(): + array = ak.Array([[0.0, 1.1, 2.2], [], [3.3, 4.4], [5.5], [6.6, 7.7, 8.8, 9.9]]) + + assert ak.combinations(array, 2, axis=-1).tolist() == [ + [(0.0, 1.1), (0.0, 2.2), (1.1, 2.2)], + [], + [(3.3, 4.4)], + [], + [(6.6, 7.7), (6.6, 8.8), (6.6, 9.9), (7.7, 8.8), (7.7, 9.9), (8.8, 9.9)], + ] + assert ak.argcombinations(array, 2, axis=-1).tolist() == [ + [(0, 1), (0, 2), (1, 2)], + [], + [(0, 1)], + [], + [(0, 1), (0, 2), (0, 3), (1, 2), (1, 3), (2, 3)], + ]