From c96f61e55e121c13fb6e262b2f198d9237b1f1d2 Mon Sep 17 00:00:00 2001 From: Peter Fackeldey Date: Fri, 13 Sep 2024 15:03:07 -0400 Subject: [PATCH] next batch of high-level functions --- src/awkward/_layout.py | 2 +- src/awkward/_namedaxis.py | 50 ++++ src/awkward/operations/ak_all.py | 5 +- src/awkward/operations/ak_any.py | 5 +- src/awkward/operations/ak_argcartesian.py | 8 +- src/awkward/operations/ak_argcombinations.py | 8 +- src/awkward/operations/ak_argmax.py | 5 +- src/awkward/operations/ak_argmin.py | 7 +- src/awkward/operations/ak_argsort.py | 7 +- src/awkward/operations/ak_categories.py | 16 +- src/awkward/operations/ak_combinations.py | 8 +- src/awkward/operations/ak_concatenate.py | 8 +- src/awkward/operations/ak_corr.py | 18 +- src/awkward/operations/ak_count.py | 47 ++- src/awkward/operations/ak_count_nonzero.py | 45 ++- src/awkward/operations/ak_covar.py | 18 +- src/awkward/operations/ak_drop_none.py | 24 +- src/awkward/operations/ak_fill_none.py | 27 +- src/awkward/operations/ak_firsts.py | 42 ++- src/awkward/operations/ak_flatten.py | 10 +- src/awkward/operations/ak_is_none.py | 26 +- src/awkward/operations/ak_linear_fit.py | 9 - src/awkward/operations/ak_local_index.py | 45 ++- src/awkward/operations/ak_max.py | 7 +- src/awkward/operations/ak_mean.py | 62 ++-- .../operations/ak_merge_option_of_records.py | 8 +- .../operations/ak_merge_union_of_records.py | 8 +- src/awkward/operations/ak_min.py | 42 ++- src/awkward/operations/ak_moment.py | 34 +-- src/awkward/operations/ak_nan_to_none.py | 9 +- src/awkward/operations/ak_nan_to_num.py | 9 +- src/awkward/operations/ak_num.py | 27 +- src/awkward/operations/ak_pad_none.py | 9 +- src/awkward/operations/ak_prod.py | 9 +- src/awkward/operations/ak_ptp.py | 9 +- src/awkward/operations/ak_real.py | 8 +- src/awkward/operations/ak_round.py | 8 +- src/awkward/operations/ak_run_lengths.py | 8 +- src/awkward/operations/ak_singletons.py | 9 +- src/awkward/operations/ak_softmax.py | 9 +- src/awkward/operations/ak_sort.py | 11 +- src/awkward/operations/ak_std.py | 9 +- src/awkward/operations/ak_strings_astype.py | 8 +- src/awkward/operations/ak_sum.py | 48 ++- src/awkward/operations/ak_transform.py | 12 +- src/awkward/operations/ak_unflatten.py | 43 ++- src/awkward/operations/ak_unzip.py | 14 +- src/awkward/operations/ak_values_astype.py | 9 +- src/awkward/operations/ak_var.py | 9 +- src/awkward/operations/ak_with_field.py | 8 +- src/awkward/operations/ak_with_name.py | 8 +- src/awkward/operations/ak_with_named_axis.py | 7 +- src/awkward/operations/ak_with_parameter.py | 8 +- src/awkward/operations/ak_without_field.py | 8 +- .../operations/ak_without_parameters.py | 8 +- src/awkward/operations/ak_zip.py | 8 +- tests/test_2596_named_axis.py | 276 +++++++++++++++++- 57 files changed, 738 insertions(+), 471 deletions(-) diff --git a/src/awkward/_layout.py b/src/awkward/_layout.py index 32ca70f101..eb6957e1e6 100644 --- a/src/awkward/_layout.py +++ b/src/awkward/_layout.py @@ -234,7 +234,7 @@ def maybe_highlevel_to_lowlevel(obj): Args: obj: an object - Calls #ak.to_layout and returns the result iff. the object is a high-level + Calls #ak.to_layout and returns the result if the object is a high-level Awkward object, otherwise the object is returned as-is. This function should be removed once scalars are properly handled by `to_layout`. diff --git a/src/awkward/_namedaxis.py b/src/awkward/_namedaxis.py index 7bd735abc2..c157948342 100644 --- a/src/awkward/_namedaxis.py +++ b/src/awkward/_namedaxis.py @@ -222,6 +222,8 @@ def _set_named_axis_to_attrs( # # The strategies are: # - "keep all" (_identity_named_axis): Keep all named axes in the output array, e.g.: `ak.drop_none` +# - "keep one" (_keep_named_axis): Keep one named axes in the output array, e.g.: `ak.firsts` +# - "remove all" (_remove_all_named_axis): Removes all named axis, e.g.: `ak.categories # - "remove one" (_remove_named_axis): Remove the named axis from the output array, e.g.: `ak.sum` # - "unify" (_unify_named_axis): Unify the named axis in the output array given two input arrays, e.g.: `__add__` # - "collapse" (_collapse_named_axis): Collapse multiple named axis to None in the output array, e.g.: `ak.flatten` @@ -250,6 +252,54 @@ def _identity_named_axis( return tuple(named_axis) +def _keep_named_axis( + named_axis: AxisTuple, + axis: int | None = None, +) -> AxisTuple: + """ + Determines the new named axis after keeping the specified axis. This is useful, for example, + when applying an operation that keeps only one axis. + + Args: + named_axis (AxisTuple): The current named axis. + axis (int | None, optional): The index of the axis to keep. If None, all axes are kept. Default is None. + + Returns: + AxisTuple: The new named axis after keeping the specified axis. + + Examples: + >>> _keep_named_axis(("x", "y", "z"), 1) + ("y",) + >>> _keep_named_axis(("x", "y", "z")) + ("x", "y", "z") + """ + return tuple(named_axis) if axis is None else (named_axis[axis],) + + +def _remove_all_named_axis( + named_axis: AxisTuple, + n: int | None = None, +) -> AxisTuple: + """ + Determines the new named axis after removing all axes. This is useful, for example, + when applying an operation that removes all axes. + + Args: + named_axis (AxisTuple): The current named axis. + n (int | None, optional): The number of axes to remove. If None, all axes are removed. Default is None. + + Returns: + AxisTuple: The new named axis after removing all axes. All elements will be None. + + Examples: + >>> _remove_all_named_axis(("x", "y", "z")) + (None, None, None) + >>> _remove_all_named_axis(("x", "y", "z"), 2) + (None, None) + """ + return (None,) * (len(named_axis) if n is None else n) + + def _remove_named_axis( axis: int | None, named_axis: AxisTuple, diff --git a/src/awkward/operations/ak_all.py b/src/awkward/operations/ak_all.py index d383df6bf1..f72f0a7e9c 100644 --- a/src/awkward/operations/ak_all.py +++ b/src/awkward/operations/ak_all.py @@ -14,6 +14,7 @@ _supports_named_axis, ) from awkward._nplikes.numpy_like import NumpyMetadata +from awkward._regularize import is_integer, regularize_axis __all__ = ("all",) @@ -91,7 +92,9 @@ def _impl(array, axis, keepdims, mask_identity, highlevel, behavior, attrs): if not keepdims: out_named_axis = _remove_named_axis(axis, out_named_axis) - if not isinstance(axis, int) and axis is not None: + axis = regularize_axis(axis) + + if not is_integer(axis) and axis is not None: raise TypeError(f"'axis' must be an integer or None by now, not {axis!r}") reducer = ak._reducers.All() diff --git a/src/awkward/operations/ak_any.py b/src/awkward/operations/ak_any.py index 2a92d5eaab..418ade73ab 100644 --- a/src/awkward/operations/ak_any.py +++ b/src/awkward/operations/ak_any.py @@ -14,6 +14,7 @@ _supports_named_axis, ) from awkward._nplikes.numpy_like import NumpyMetadata +from awkward._regularize import is_integer, regularize_axis __all__ = ("any",) @@ -91,7 +92,9 @@ def _impl(array, axis, keepdims, mask_identity, highlevel, behavior, attrs): if not keepdims: out_named_axis = _remove_named_axis(axis, out_named_axis) - if not isinstance(axis, int) and axis is not None: + axis = regularize_axis(axis) + + if not is_integer(axis) and axis is not None: raise TypeError(f"'axis' must be an integer or None by now, not {axis!r}") reducer = ak._reducers.Any() diff --git a/src/awkward/operations/ak_argcartesian.py b/src/awkward/operations/ak_argcartesian.py index 50fea77736..12deed5749 100644 --- a/src/awkward/operations/ak_argcartesian.py +++ b/src/awkward/operations/ak_argcartesian.py @@ -6,9 +6,8 @@ import awkward as ak from awkward._dispatch import high_level_function -from awkward._namedaxis import _supports_named_axis from awkward._nplikes.numpy_like import NumpyMetadata -from awkward._regularize import is_integer, regularize_axis +from awkward._regularize import regularize_axis __all__ = ("argcartesian",) @@ -108,11 +107,6 @@ def argcartesian( def _impl(arrays, axis, nested, parameters, with_name, highlevel, behavior, attrs): - out_named_axis = None - if _supports_named_axis(arrays[0]) and not is_integer(axis): - # Named axis handling - raise NotImplementedError() - axis = regularize_axis(axis) if isinstance(arrays, Mapping): diff --git a/src/awkward/operations/ak_argcombinations.py b/src/awkward/operations/ak_argcombinations.py index 04fd3fd158..c2b4793c95 100644 --- a/src/awkward/operations/ak_argcombinations.py +++ b/src/awkward/operations/ak_argcombinations.py @@ -5,9 +5,8 @@ import awkward as ak from awkward._dispatch import high_level_function from awkward._layout import HighLevelContext -from awkward._namedaxis import _supports_named_axis from awkward._nplikes.numpy_like import NumpyMetadata -from awkward._regularize import is_integer, regularize_axis +from awkward._regularize import regularize_axis __all__ = ("argcombinations",) @@ -94,11 +93,6 @@ def _impl( behavior, attrs, ): - out_named_axis = None - if _supports_named_axis(array) and not is_integer(axis): - # Named axis handling - raise NotImplementedError() - axis = regularize_axis(axis) if parameters is None: diff --git a/src/awkward/operations/ak_argmax.py b/src/awkward/operations/ak_argmax.py index 9691de1ecc..c487eb5b8c 100644 --- a/src/awkward/operations/ak_argmax.py +++ b/src/awkward/operations/ak_argmax.py @@ -14,6 +14,7 @@ _supports_named_axis, ) from awkward._nplikes.numpy_like import NumpyMetadata +from awkward._regularize import is_integer, regularize_axis __all__ = ("argmax", "nanargmax") @@ -156,7 +157,9 @@ def _impl(array, axis, keepdims, mask_identity, highlevel, behavior, attrs): if not keepdims: out_named_axis = _remove_named_axis(axis, out_named_axis) - if not isinstance(axis, int) and axis is not None: + axis = regularize_axis(axis) + + if not is_integer(axis) and axis is not None: raise TypeError(f"'axis' must be an integer or None by now, not {axis!r}") # axis = regularize_axis(axis) diff --git a/src/awkward/operations/ak_argmin.py b/src/awkward/operations/ak_argmin.py index 1c1934f047..d0cea0b9c2 100644 --- a/src/awkward/operations/ak_argmin.py +++ b/src/awkward/operations/ak_argmin.py @@ -14,6 +14,7 @@ _supports_named_axis, ) from awkward._nplikes.numpy_like import NumpyMetadata +from awkward._regularize import is_integer, regularize_axis __all__ = ("argmin", "nanargmin") @@ -153,10 +154,10 @@ def _impl(array, axis, keepdims, mask_identity, highlevel, behavior, attrs): if not keepdims: out_named_axis = _remove_named_axis(axis, out_named_axis) - if not isinstance(axis, int) and axis is not None: - raise TypeError(f"'axis' must be an integer or None by now, not {axis!r}") + axis = regularize_axis(axis) - # axis = regularize_axis(axis) + if not is_integer(axis) and axis is not None: + raise TypeError(f"'axis' must be an integer or None by now, not {axis!r}") reducer = ak._reducers.ArgMin() diff --git a/src/awkward/operations/ak_argsort.py b/src/awkward/operations/ak_argsort.py index 58bb22875a..fbc6223607 100644 --- a/src/awkward/operations/ak_argsort.py +++ b/src/awkward/operations/ak_argsort.py @@ -13,6 +13,7 @@ _supports_named_axis, ) from awkward._nplikes.numpy_like import NumpyMetadata +from awkward._regularize import is_integer, regularize_axis __all__ = ("argsort",) @@ -90,8 +91,10 @@ def _impl(array, axis, ascending, stable, highlevel, behavior, attrs): # use strategy "keep all" (see: awkward._namedaxis) out_named_axis = _identity_named_axis(array.named_axis) - if not isinstance(axis, int) and axis is not None: - raise TypeError(f"'axis' must be an integer or None by now, not {axis!r}") + axis = regularize_axis(axis) + + if not is_integer(axis): + raise TypeError(f"'axis' must be an integer by now, not {axis!r}") out = ak._do.argsort(layout, axis, ascending, stable) diff --git a/src/awkward/operations/ak_categories.py b/src/awkward/operations/ak_categories.py index cd7f6ccf4c..38cb4c87ed 100644 --- a/src/awkward/operations/ak_categories.py +++ b/src/awkward/operations/ak_categories.py @@ -5,6 +5,7 @@ import awkward as ak from awkward._dispatch import high_level_function from awkward._layout import HighLevelContext +from awkward._namedaxis import _remove_all_named_axis __all__ = ("categories",) @@ -49,6 +50,19 @@ def action(layout, **kwargs): with HighLevelContext(behavior=behavior, attrs=attrs) as ctx: layout = ctx.unwrap(array, allow_record=False, primitive_policy="error") + ak._do.recursively_apply(layout, action) - return ctx.wrap(output, highlevel=highlevel) + wrapped_out = ctx.wrap(output, highlevel=highlevel) + + # propagate named axis from input to output, + # use strategy "drop all" (see: awkward._namedaxis) + out_named_axis = _remove_all_named_axis(wrapped_out.named_axis, n=wrapped_out.ndim) + + return ak.operations.ak_with_named_axis._impl( + wrapped_out, + named_axis=out_named_axis, + highlevel=highlevel, + behavior=ctx.behavior, + attrs=ctx.attrs, + ) diff --git a/src/awkward/operations/ak_combinations.py b/src/awkward/operations/ak_combinations.py index a7e0fba1bf..d22708cb4a 100644 --- a/src/awkward/operations/ak_combinations.py +++ b/src/awkward/operations/ak_combinations.py @@ -5,9 +5,8 @@ import awkward as ak from awkward._dispatch import high_level_function from awkward._layout import HighLevelContext -from awkward._namedaxis import _supports_named_axis from awkward._nplikes.numpy_like import NumpyMetadata -from awkward._regularize import is_integer, regularize_axis +from awkward._regularize import regularize_axis __all__ = ("combinations",) @@ -215,11 +214,6 @@ def _impl( behavior, attrs, ): - out_named_axis = None - if _supports_named_axis(array) and not is_integer(axis): - # Named axis handling - raise NotImplementedError() - axis = regularize_axis(axis) if with_name is None: diff --git a/src/awkward/operations/ak_concatenate.py b/src/awkward/operations/ak_concatenate.py index accbb0c9a9..84a38717a9 100644 --- a/src/awkward/operations/ak_concatenate.py +++ b/src/awkward/operations/ak_concatenate.py @@ -9,11 +9,10 @@ from awkward._dispatch import high_level_function from awkward._do import mergeable from awkward._layout import HighLevelContext, ensure_same_backend, maybe_posaxis -from awkward._namedaxis import _supports_named_axis from awkward._nplikes.numpy_like import NumpyMetadata from awkward._nplikes.shape import unknown_length from awkward._parameters import type_parameters_equal -from awkward._regularize import is_integer, regularize_axis +from awkward._regularize import regularize_axis from awkward._typing import Sequence from awkward.contents import Content from awkward.operations.ak_fill_none import fill_none @@ -93,11 +92,6 @@ def _merge_as_union( def _impl(arrays, axis, mergebool, highlevel, behavior, attrs): - out_named_axis = None - if _supports_named_axis(arrays[0]) and not is_integer(axis): - # Named axis handling - raise NotImplementedError() - axis = regularize_axis(axis) # Simple single-array, axis=0 fast-path diff --git a/src/awkward/operations/ak_corr.py b/src/awkward/operations/ak_corr.py index 585b7bc712..9300bda77c 100644 --- a/src/awkward/operations/ak_corr.py +++ b/src/awkward/operations/ak_corr.py @@ -9,10 +9,8 @@ ensure_same_backend, maybe_highlevel_to_lowlevel, ) -from awkward._namedaxis import _supports_named_axis from awkward._nplikes import ufuncs from awkward._nplikes.numpy_like import NumpyMetadata -from awkward._regularize import is_integer, regularize_axis __all__ = ("corr",) @@ -87,13 +85,6 @@ def corr( def _impl(x, y, weight, axis, keepdims, mask_identity, highlevel, behavior, attrs): - out_named_axis = None - if _supports_named_axis(x) and not is_integer(axis): - # Named axis handling - raise NotImplementedError() - - axis = regularize_axis(axis) - with HighLevelContext(behavior=behavior, attrs=attrs) as ctx: x_layout, y_layout, weight_layout = ensure_same_backend( ctx.unwrap(x, allow_record=False, primitive_policy="error"), @@ -190,8 +181,13 @@ def _impl(x, y, weight, axis, keepdims, mask_identity, highlevel, behavior, attr behavior=ctx.behavior, attrs=ctx.attrs, ) - return ctx.wrap( - maybe_highlevel_to_lowlevel(sumwxy / ufuncs.sqrt(sumwxx * sumwyy)), + + # propagate named axis to output + out = sumwxy / ufuncs.sqrt(sumwxx * sumwyy) + out_ctx = HighLevelContext(behavior=out.behavior, attrs=out.attrs).finalize() + + return out_ctx.wrap( + maybe_highlevel_to_lowlevel(out), highlevel=highlevel, allow_other=True, ) diff --git a/src/awkward/operations/ak_count.py b/src/awkward/operations/ak_count.py index 198761a992..a685c45b01 100644 --- a/src/awkward/operations/ak_count.py +++ b/src/awkward/operations/ak_count.py @@ -5,7 +5,13 @@ import awkward as ak from awkward._dispatch import high_level_function from awkward._layout import HighLevelContext -from awkward._namedaxis import _supports_named_axis +from awkward._namedaxis import ( + _check_valid_axis, + _identity_named_axis, + _one_axis_to_positional_axis, + _remove_named_axis, + _supports_named_axis, +) from awkward._nplikes.numpy_like import NumpyMetadata from awkward._regularize import is_integer, regularize_axis @@ -110,15 +116,29 @@ def count( def _impl(array, axis, keepdims, mask_identity, highlevel, behavior, attrs): + with HighLevelContext(behavior=behavior, attrs=attrs) as ctx: + layout = ctx.unwrap(array, allow_record=False, primitive_policy="error") + out_named_axis = None - if _supports_named_axis(array) and not is_integer(axis): - # Named axis handling - raise NotImplementedError() + if _supports_named_axis(ctx) and _check_valid_axis(axis): + # Handle named axis + # Step 1: Normalize named axis to positional axis + axis = _one_axis_to_positional_axis( + axis, array.named_axis, array.positional_axis + ) + + # Step 2: propagate named axis from input to output, + # keepdims=True: use strategy "keep all" (see: awkward._namedaxis) + # keepdims=False: use strategy "remove one" (see: awkward._namedaxis) + out_named_axis = _identity_named_axis(array.named_axis) + if not keepdims: + out_named_axis = _remove_named_axis(axis, out_named_axis) axis = regularize_axis(axis) - with HighLevelContext(behavior=behavior, attrs=attrs) as ctx: - layout = ctx.unwrap(array, allow_record=False, primitive_policy="error") + if not is_integer(axis) and axis is not None: + raise TypeError(f"'axis' must be an integer or None by now, not {axis!r}") + reducer = ak._reducers.Count() out = ak._do.reduce( @@ -129,9 +149,20 @@ def _impl(array, axis, keepdims, mask_identity, highlevel, behavior, attrs): keepdims=keepdims, behavior=ctx.behavior, ) - return ctx.wrap( + + wrapped_out = ctx.wrap( out, highlevel=highlevel, allow_other=True, - named_axis=out_named_axis, ) + + if out_named_axis: + # propagate named axis to output + return ak.operations.ak_with_named_axis._impl( + wrapped_out, + named_axis=out_named_axis, + highlevel=highlevel, + behavior=ctx.behavior, + attrs=ctx.attrs, + ) + return wrapped_out diff --git a/src/awkward/operations/ak_count_nonzero.py b/src/awkward/operations/ak_count_nonzero.py index 8e7a4be5cd..8cae937e11 100644 --- a/src/awkward/operations/ak_count_nonzero.py +++ b/src/awkward/operations/ak_count_nonzero.py @@ -5,7 +5,13 @@ import awkward as ak from awkward._dispatch import high_level_function from awkward._layout import HighLevelContext -from awkward._namedaxis import _supports_named_axis +from awkward._namedaxis import ( + _check_valid_axis, + _identity_named_axis, + _one_axis_to_positional_axis, + _remove_named_axis, + _supports_named_axis, +) from awkward._nplikes.numpy_like import NumpyMetadata from awkward._regularize import is_integer, regularize_axis @@ -69,13 +75,29 @@ def count_nonzero( def _impl(array, axis, keepdims, mask_identity, highlevel, behavior, attrs): + with HighLevelContext(behavior=behavior, attrs=attrs) as ctx: + layout = ctx.unwrap(array, allow_record=False, primitive_policy="error") + out_named_axis = None - if _supports_named_axis(array) and not is_integer(axis): - # Named axis handling - raise NotImplementedError() + if _supports_named_axis(ctx) and _check_valid_axis(axis): + # Handle named axis + # Step 1: Normalize named axis to positional axis + axis = _one_axis_to_positional_axis( + axis, array.named_axis, array.positional_axis + ) + + # Step 2: propagate named axis from input to output, + # keepdims=True: use strategy "keep all" (see: awkward._namedaxis) + # keepdims=False: use strategy "remove one" (see: awkward._namedaxis) + out_named_axis = _identity_named_axis(array.named_axis) + if not keepdims: + out_named_axis = _remove_named_axis(axis, out_named_axis) axis = regularize_axis(axis) + if not is_integer(axis) and axis is not None: + raise TypeError(f"'axis' must be an integer or None by now, not {axis!r}") + with HighLevelContext(behavior=behavior, attrs=attrs) as ctx: layout = ctx.unwrap(array, allow_record=False, primitive_policy="error") reducer = ak._reducers.CountNonzero() @@ -88,13 +110,24 @@ def _impl(array, axis, keepdims, mask_identity, highlevel, behavior, attrs): keepdims=keepdims, behavior=ctx.behavior, ) - return ctx.wrap( + + wrapped_out = ctx.wrap( out, highlevel=highlevel, allow_other=True, - named_axis=out_named_axis, ) + if out_named_axis: + # propagate named axis to output + return ak.operations.ak_with_named_axis._impl( + wrapped_out, + named_axis=out_named_axis, + highlevel=highlevel, + behavior=ctx.behavior, + attrs=ctx.attrs, + ) + return wrapped_out + @ak._connect.numpy.implements("count_nonzero") def _nep_18_impl(a, axis=None, *, keepdims=False): diff --git a/src/awkward/operations/ak_covar.py b/src/awkward/operations/ak_covar.py index 74bfc484a1..e8dcf5f844 100644 --- a/src/awkward/operations/ak_covar.py +++ b/src/awkward/operations/ak_covar.py @@ -9,9 +9,7 @@ ensure_same_backend, maybe_highlevel_to_lowlevel, ) -from awkward._namedaxis import _supports_named_axis from awkward._nplikes.numpy_like import NumpyMetadata -from awkward._regularize import is_integer, regularize_axis __all__ = ("covar",) @@ -84,13 +82,6 @@ def covar( def _impl(x, y, weight, axis, keepdims, mask_identity, highlevel, behavior, attrs): - out_named_axis = None - if _supports_named_axis(x) and not is_integer(axis): - # Named axis handling - raise NotImplementedError() - - axis = regularize_axis(axis) - with HighLevelContext(behavior=behavior, attrs=attrs) as ctx: x_layout, y_layout, weight_layout = ensure_same_backend( ctx.unwrap(x, allow_record=False, primitive_policy="error"), @@ -167,8 +158,13 @@ def _impl(x, y, weight, axis, keepdims, mask_identity, highlevel, behavior, attr behavior=None, attrs=None, ) - return ctx.wrap( - maybe_highlevel_to_lowlevel(sumwxy / sumw), + + # propagate named axis to output + out = sumwxy / sumw + out_ctx = HighLevelContext(behavior=out.behavior, attrs=out.attrs).finalize() + + return out_ctx.wrap( + maybe_highlevel_to_lowlevel(out), highlevel=highlevel, allow_other=True, ) diff --git a/src/awkward/operations/ak_drop_none.py b/src/awkward/operations/ak_drop_none.py index c568ef2fe7..596a986073 100644 --- a/src/awkward/operations/ak_drop_none.py +++ b/src/awkward/operations/ak_drop_none.py @@ -5,7 +5,11 @@ import awkward as ak from awkward._dispatch import high_level_function from awkward._layout import HighLevelContext, maybe_posaxis -from awkward._namedaxis import _supports_named_axis +from awkward._namedaxis import ( + _check_valid_axis, + _one_axis_to_positional_axis, + _supports_named_axis, +) from awkward._nplikes.numpy_like import NumpyMetadata from awkward._regularize import is_integer, regularize_axis from awkward.errors import AxisError @@ -66,15 +70,20 @@ def _drop_none_if_list(layout): def _impl(array, axis, highlevel, behavior, attrs): - out_named_axis = None - if _supports_named_axis(array) and not is_integer(axis): - # Named axis handling - raise NotImplementedError() + with HighLevelContext(behavior=behavior, attrs=attrs) as ctx: + layout = ctx.unwrap(array, allow_record=False, primitive_policy="error") + + if _supports_named_axis(ctx) and _check_valid_axis(axis): + # Handle named axis + # Step 1: Normalize named axis to positional axis + axis = _one_axis_to_positional_axis( + axis, array.named_axis, array.positional_axis + ) axis = regularize_axis(axis) - with HighLevelContext(behavior=behavior, attrs=attrs) as ctx: - layout = ctx.unwrap(array, allow_record=False, primitive_policy="error") + if not is_integer(axis) and axis is not None: + raise TypeError(f"'axis' must be an integer or None by now, not {axis!r}") if axis is None: # if the outer layout is_option, drop_nones without affecting offsets @@ -130,5 +139,4 @@ def action(layout, depth, **kwargs): return ctx.wrap( out, highlevel=highlevel, - named_axis=out_named_axis, ) diff --git a/src/awkward/operations/ak_fill_none.py b/src/awkward/operations/ak_fill_none.py index 3de222679b..a365503c86 100644 --- a/src/awkward/operations/ak_fill_none.py +++ b/src/awkward/operations/ak_fill_none.py @@ -5,7 +5,11 @@ import awkward as ak from awkward._dispatch import high_level_function from awkward._layout import HighLevelContext, ensure_same_backend, maybe_posaxis -from awkward._namedaxis import _supports_named_axis +from awkward._namedaxis import ( + _check_valid_axis, + _one_axis_to_positional_axis, + _supports_named_axis, +) from awkward._nplikes.numpy_like import NumpyMetadata from awkward._regularize import is_integer, regularize_axis from awkward.errors import AxisError @@ -70,13 +74,6 @@ def fill_none(array, value, axis=-1, *, highlevel=True, behavior=None, attrs=Non def _impl(array, value, axis, highlevel, behavior, attrs): - out_named_axis = None - if _supports_named_axis(array) and not is_integer(axis): - # Named axis handling - raise NotImplementedError() - - axis = regularize_axis(axis) - with HighLevelContext(behavior=behavior, attrs=attrs) as ctx: array_layout, value_layout = ensure_same_backend( ctx.unwrap(array, allow_record=True, allow_unknown=False), @@ -90,6 +87,18 @@ def _impl(array, value, axis, highlevel, behavior, attrs): ), ) + if _supports_named_axis(ctx) and _check_valid_axis(axis): + # Handle named axis + # Step 1: Normalize named axis to positional axis + axis = _one_axis_to_positional_axis( + axis, array.named_axis, array.positional_axis + ) + + axis = regularize_axis(axis) + + if not is_integer(axis) and axis is not None: + raise TypeError(f"'axis' must be an integer or None by now, not {axis!r}") + if isinstance(value_layout, ak.record.Record): value_layout = value_layout.array[value_layout.at : value_layout.at + 1] elif isinstance(value_layout, ak.contents.Content): @@ -131,4 +140,4 @@ def action(layout, depth, **kwargs): ) out = ak._do.recursively_apply(array_layout, action) - return ctx.wrap(out, highlevel=highlevel, named_axis=out_named_axis) + return ctx.wrap(out, highlevel=highlevel) diff --git a/src/awkward/operations/ak_firsts.py b/src/awkward/operations/ak_firsts.py index 847db640fb..3b87620bc0 100644 --- a/src/awkward/operations/ak_firsts.py +++ b/src/awkward/operations/ak_firsts.py @@ -5,7 +5,12 @@ import awkward as ak from awkward._dispatch import high_level_function from awkward._layout import HighLevelContext, maybe_posaxis -from awkward._namedaxis import _supports_named_axis +from awkward._namedaxis import ( + _check_valid_axis, + _keep_named_axis, + _one_axis_to_positional_axis, + _supports_named_axis, +) from awkward._nplikes.numpy_like import NumpyMetadata from awkward._regularize import is_integer, regularize_axis from awkward.errors import AxisError @@ -57,18 +62,25 @@ def firsts(array, axis=1, *, highlevel=True, behavior=None, attrs=None): def _impl(array, axis, highlevel, behavior, attrs): + with HighLevelContext(behavior=behavior, attrs=attrs) as ctx: + layout = ctx.unwrap(array, allow_record=False) + out_named_axis = None - if _supports_named_axis(array) and not is_integer(axis): - # Named axis handling - raise NotImplementedError() + if _supports_named_axis(ctx) and _check_valid_axis(axis): + # Handle named axis + # Step 1: Normalize named axis to positional axis + axis = _one_axis_to_positional_axis( + axis, array.named_axis, array.positional_axis + ) - axis = regularize_axis(axis) + # Step 2: propagate named axis from input to output, + # use strategy "keep one" (see: awkward._namedaxis) + out_named_axis = _keep_named_axis(array.named_axis, axis) - with HighLevelContext(behavior=behavior, attrs=attrs) as ctx: - layout = ctx.unwrap(array, allow_record=False) + axis = regularize_axis(axis) if not is_integer(axis): - raise TypeError(f"'axis' must be an integer, not {axis!r}") + raise TypeError(f"'axis' must be an integer by now, not {axis!r}") if maybe_posaxis(layout, axis, 1) == 0: # specialized logic; it's tested in test_0582-propagate-context-in-broadcast_and_apply.py @@ -110,9 +122,19 @@ def action(layout, depth, backend, **kwargs): out = ak._do.recursively_apply(layout, action, numpy_to_regular=True) - return ctx.wrap( + wrapped_out = ctx.wrap( out, highlevel=highlevel, allow_other=True, - named_axis=out_named_axis, ) + + if out_named_axis: + # propagate named axis to output + return ak.operations.ak_with_named_axis._impl( + wrapped_out, + named_axis=out_named_axis, + highlevel=highlevel, + behavior=ctx.behavior, + attrs=ctx.attrs, + ) + return wrapped_out diff --git a/src/awkward/operations/ak_flatten.py b/src/awkward/operations/ak_flatten.py index 08fb5d6c70..6668a06bdd 100644 --- a/src/awkward/operations/ak_flatten.py +++ b/src/awkward/operations/ak_flatten.py @@ -5,9 +5,8 @@ import awkward as ak from awkward._dispatch import high_level_function from awkward._layout import HighLevelContext, maybe_posaxis -from awkward._namedaxis import _supports_named_axis from awkward._nplikes.numpy_like import NumpyMetadata -from awkward._regularize import is_integer, regularize_axis +from awkward._regularize import regularize_axis __all__ = ("flatten",) @@ -174,11 +173,6 @@ def flatten(array, axis=1, *, highlevel=True, behavior=None, attrs=None): def _impl(array, axis, highlevel, behavior, attrs): - out_named_axis = None - if _supports_named_axis(array) and not is_integer(axis): - # Named axis handling - raise NotImplementedError() - axis = regularize_axis(axis) with HighLevelContext(behavior=behavior, attrs=attrs) as ctx: @@ -241,4 +235,4 @@ def apply(layout): out = apply(layout) else: out = ak._do.flatten(layout, axis) - return ctx.wrap(out, highlevel=highlevel, named_axis=out_named_axis) + return ctx.wrap(out, highlevel=highlevel) diff --git a/src/awkward/operations/ak_is_none.py b/src/awkward/operations/ak_is_none.py index 5e391e463f..340bf9ae3f 100644 --- a/src/awkward/operations/ak_is_none.py +++ b/src/awkward/operations/ak_is_none.py @@ -6,12 +6,12 @@ from awkward._dispatch import high_level_function from awkward._layout import HighLevelContext, maybe_posaxis from awkward._namedaxis import ( - _identity_named_axis, + _check_valid_axis, _one_axis_to_positional_axis, _supports_named_axis, ) from awkward._nplikes.numpy_like import NumpyMetadata -from awkward._regularize import is_integer +from awkward._regularize import is_integer, regularize_axis from awkward.errors import AxisError __all__ = ("is_none",) @@ -46,28 +46,20 @@ def is_none(array, axis=0, *, highlevel=True, behavior=None, attrs=None): def _impl(array, axis, highlevel, behavior, attrs): - out_named_axis = None - if _supports_named_axis(array) and not is_integer(axis): + with HighLevelContext(behavior=behavior, attrs=attrs) as ctx: + layout = ctx.unwrap(array, allow_record=False, primitive_policy="error") + + if _supports_named_axis(ctx) and _check_valid_axis(axis): # Handle named axis # Step 1: Normalize named axis to positional axis axis = _one_axis_to_positional_axis( axis, array.named_axis, array.positional_axis ) - # Step 2: propagate named axis from input to output, - # use strategy "remove one" (see: awkward._namedaxis) - out_named_axis = _identity_named_axis(array.named_axis) - - if not isinstance(axis, int): - raise TypeError(f"'axis' must be an integer by now, not {axis!r}") - - # axis = regularize_axis(axis) # <- is this really needed? - - with HighLevelContext(behavior=behavior, attrs=attrs) as ctx: - layout = ctx.unwrap(array, allow_record=False, primitive_policy="error") + axis = regularize_axis(axis) if not is_integer(axis): - raise TypeError(f"'axis' must be an integer, not {axis!r}") + raise TypeError(f"'axis' must be an integer by now, not {axis!r}") def action(layout, depth, backend, lateral_context, **kwargs): posaxis = maybe_posaxis(layout, axis, depth) @@ -89,4 +81,4 @@ def action(layout, depth, backend, lateral_context, **kwargs): out = ak._do.recursively_apply(layout, action, numpy_to_regular=True) - return ctx.wrap(out, highlevel=highlevel, named_axis=out_named_axis) + return ctx.wrap(out, highlevel=highlevel) diff --git a/src/awkward/operations/ak_linear_fit.py b/src/awkward/operations/ak_linear_fit.py index d9e17e14eb..14c21019de 100644 --- a/src/awkward/operations/ak_linear_fit.py +++ b/src/awkward/operations/ak_linear_fit.py @@ -5,10 +5,8 @@ import awkward as ak from awkward._dispatch import high_level_function from awkward._layout import HighLevelContext, ensure_same_backend -from awkward._namedaxis import _supports_named_axis from awkward._nplikes import ufuncs from awkward._nplikes.numpy_like import NumpyMetadata -from awkward._regularize import is_integer, regularize_axis __all__ = ("linear_fit",) @@ -96,13 +94,6 @@ def linear_fit( def _impl(x, y, weight, axis, keepdims, mask_identity, highlevel, behavior, attrs): - out_named_axis = None - if _supports_named_axis(x) and not is_integer(axis): - # Named axis handling - raise NotImplementedError() - - axis = regularize_axis(axis) - with HighLevelContext(behavior=behavior, attrs=attrs) as ctx: x_layout, y_layout, weight_layout = ensure_same_backend( ctx.unwrap(x, allow_record=False, primitive_policy="error"), diff --git a/src/awkward/operations/ak_local_index.py b/src/awkward/operations/ak_local_index.py index 3817166451..d8aa911c68 100644 --- a/src/awkward/operations/ak_local_index.py +++ b/src/awkward/operations/ak_local_index.py @@ -5,7 +5,12 @@ import awkward as ak from awkward._dispatch import high_level_function from awkward._layout import HighLevelContext -from awkward._namedaxis import _supports_named_axis +from awkward._namedaxis import ( + _check_valid_axis, + _keep_named_axis, + _one_axis_to_positional_axis, + _supports_named_axis, +) from awkward._nplikes.numpy_like import NumpyMetadata from awkward._regularize import is_integer, regularize_axis @@ -89,14 +94,40 @@ def local_index(array, axis=-1, *, highlevel=True, behavior=None, attrs=None): def _impl(array, axis, highlevel, behavior, attrs): + with HighLevelContext(behavior=behavior, attrs=attrs) as ctx: + layout = ctx.unwrap(array, allow_record=False, primitive_policy="error") + out_named_axis = None - if _supports_named_axis(array) and not is_integer(axis): - # Named axis handling - raise NotImplementedError() + if _supports_named_axis(ctx) and _check_valid_axis(axis): + # Handle named axis + # Step 1: Normalize named axis to positional axis + axis = _one_axis_to_positional_axis( + axis, array.named_axis, array.positional_axis + ) axis = regularize_axis(axis) - with HighLevelContext(behavior=behavior, attrs=attrs) as ctx: - layout = ctx.unwrap(array, allow_record=False, primitive_policy="error") + if not is_integer(axis): + raise TypeError(f"'axis' must be an integer by now, not {axis!r}") + + # Step 2: propagate named axis from input to output, + # "keep all" up to the positional axis dim (see: awkward._namedaxis) + out_named_axis = _keep_named_axis(array.named_axis, None)[: axis + 1] + out = ak._do.local_index(layout, axis) - return ctx.wrap(out, highlevel=highlevel) + + wrapped_out = ctx.wrap( + out, + highlevel=highlevel, + ) + + if out_named_axis: + # propagate named axis to output + return ak.operations.ak_with_named_axis._impl( + wrapped_out, + named_axis=out_named_axis, + highlevel=highlevel, + behavior=ctx.behavior, + attrs=ctx.attrs, + ) + return wrapped_out diff --git a/src/awkward/operations/ak_max.py b/src/awkward/operations/ak_max.py index bfbc4d1e8a..4d6db85f84 100644 --- a/src/awkward/operations/ak_max.py +++ b/src/awkward/operations/ak_max.py @@ -14,6 +14,7 @@ _supports_named_axis, ) from awkward._nplikes.numpy_like import NumpyMetadata +from awkward._regularize import is_integer, regularize_axis __all__ = ("max", "nanmax") @@ -166,10 +167,10 @@ def _impl(array, axis, keepdims, initial, mask_identity, highlevel, behavior, at if not keepdims: out_named_axis = _remove_named_axis(axis, out_named_axis) - if not isinstance(axis, int) and axis is not None: - raise TypeError(f"'axis' must be an integer or None by now, not {axis!r}") + axis = regularize_axis(axis) - # axis = regularize_axis(axis) + if not is_integer(axis) and axis is not None: + raise TypeError(f"'axis' must be an integer or None by now, not {axis!r}") reducer = ak._reducers.Max(initial) diff --git a/src/awkward/operations/ak_mean.py b/src/awkward/operations/ak_mean.py index b1eacff552..399a8f560b 100644 --- a/src/awkward/operations/ak_mean.py +++ b/src/awkward/operations/ak_mean.py @@ -12,13 +12,14 @@ maybe_posaxis, ) from awkward._namedaxis import ( + _check_valid_axis, _identity_named_axis, _one_axis_to_positional_axis, _remove_named_axis, _supports_named_axis, ) from awkward._nplikes.numpy_like import NumpyMetadata -from awkward._regularize import is_integer +from awkward._regularize import is_integer, regularize_axis __all__ = ("mean", "nanmean") @@ -180,33 +181,6 @@ def nanmean( def _impl(x, weight, axis, keepdims, mask_identity, highlevel, behavior, attrs): - out_named_axis = None - if _supports_named_axis(x) and not is_integer(axis): - # Handle named axis - # Step 1: Convert named axis to positional axis - xaxis = _one_axis_to_positional_axis(axis, x.named_axis, x.positional_axis) - waxis = xaxis - if weight is not None and _supports_named_axis(weight): - waxis = _one_axis_to_positional_axis( - axis, weight.named_axis, weight.positional_axis - ) - if xaxis != waxis: - raise ValueError( - f"ak.mean require the same axis for x and weight, got {xaxis} and {waxis}" - ) - axis = xaxis - - # Step 2: Propagate named axis from input to output using "remove one" strategy - out_named_axis = _identity_named_axis(x.named_axis) - if not keepdims: - # Remove the axis that we are reducing - out_named_axis = _remove_named_axis(axis, out_named_axis) - - if not isinstance(axis, int): - raise TypeError( - f"'axis' must be an integer (or a named axis that maps to an integer) by now, not {axis!r}" - ) - with HighLevelContext(behavior=behavior, attrs=attrs) as ctx: x_layout, weight_layout = ensure_same_backend( ctx.unwrap(x, allow_record=False, primitive_policy="error"), @@ -222,6 +196,24 @@ def _impl(x, weight, axis, keepdims, mask_identity, highlevel, behavior, attrs): x = ctx.wrap(x_layout) weight = ctx.wrap(weight_layout, allow_other=True) + out_named_axis = None + if _supports_named_axis(ctx) and _check_valid_axis(axis): + # Handle named axis + # Step 1: Normalize named axis to positional axis + axis = _one_axis_to_positional_axis(axis, x.named_axis, x.positional_axis) + + # Step 2: propagate named axis from input to output, + # keepdims=True: use strategy "keep all" (see: awkward._namedaxis) + # keepdims=False: use strategy "remove one" (see: awkward._namedaxis) + out_named_axis = _identity_named_axis(x.named_axis) + if not keepdims: + out_named_axis = _remove_named_axis(axis, out_named_axis) + + axis = regularize_axis(axis) + + if not is_integer(axis) and axis is not None: + raise TypeError(f"'axis' must be an integer or None by now, not {axis!r}") + with np.errstate(invalid="ignore", divide="ignore"): if weight is None: sumw = ak.operations.ak_count._impl( @@ -282,13 +274,23 @@ def _impl(x, weight, axis, keepdims, mask_identity, highlevel, behavior, attrs): posaxis = maybe_posaxis(out.layout, axis, 1) out = out[(slice(None, None),) * posaxis + (0,)] - return ctx.wrap( + wrapped_out = ctx.wrap( maybe_highlevel_to_lowlevel(out), highlevel=highlevel, allow_other=True, - named_axis=out_named_axis, ) + if out_named_axis: + # propagate named axis to output + return ak.operations.ak_with_named_axis._impl( + wrapped_out, + named_axis=out_named_axis, + highlevel=highlevel, + behavior=ctx.behavior, + attrs=ctx.attrs, + ) + return wrapped_out + @ak._connect.numpy.implements("mean") def _nep_18_impl_mean( diff --git a/src/awkward/operations/ak_merge_option_of_records.py b/src/awkward/operations/ak_merge_option_of_records.py index 429107387b..ef0c6b4227 100644 --- a/src/awkward/operations/ak_merge_option_of_records.py +++ b/src/awkward/operations/ak_merge_option_of_records.py @@ -5,9 +5,8 @@ import awkward as ak from awkward._dispatch import high_level_function from awkward._layout import HighLevelContext, maybe_posaxis -from awkward._namedaxis import _supports_named_axis from awkward._nplikes.numpy_like import NumpyMetadata -from awkward._regularize import is_integer, regularize_axis +from awkward._regularize import regularize_axis from awkward.errors import AxisError __all__ = ("merge_option_of_records",) @@ -50,11 +49,6 @@ def merge_option_of_records( def _impl(array, axis, highlevel, behavior, attrs): - out_named_axis = None - if _supports_named_axis(array) and not is_integer(axis): - # Named axis handling - raise NotImplementedError() - axis = regularize_axis(axis) with HighLevelContext(behavior=behavior, attrs=attrs) as ctx: diff --git a/src/awkward/operations/ak_merge_union_of_records.py b/src/awkward/operations/ak_merge_union_of_records.py index e68fa6d1c7..bbd066cd06 100644 --- a/src/awkward/operations/ak_merge_union_of_records.py +++ b/src/awkward/operations/ak_merge_union_of_records.py @@ -5,9 +5,8 @@ import awkward as ak from awkward._dispatch import high_level_function from awkward._layout import HighLevelContext, maybe_posaxis -from awkward._namedaxis import _supports_named_axis from awkward._nplikes.numpy_like import ArrayLike, NumpyMetadata -from awkward._regularize import is_integer, regularize_axis +from awkward._regularize import regularize_axis from awkward.errors import AxisError __all__ = ("merge_union_of_records",) @@ -60,11 +59,6 @@ def merge_union_of_records( def _impl(array, axis, highlevel, behavior, attrs): - out_named_axis = None - if _supports_named_axis(array) and not is_integer(axis): - # Named axis handling - raise NotImplementedError() - axis = regularize_axis(axis) with HighLevelContext(behavior=behavior, attrs=attrs) as ctx: diff --git a/src/awkward/operations/ak_min.py b/src/awkward/operations/ak_min.py index bebbc851e2..a21bf94153 100644 --- a/src/awkward/operations/ak_min.py +++ b/src/awkward/operations/ak_min.py @@ -7,13 +7,14 @@ from awkward._dispatch import high_level_function from awkward._layout import HighLevelContext from awkward._namedaxis import ( + _check_valid_axis, _identity_named_axis, _one_axis_to_positional_axis, _remove_named_axis, _supports_named_axis, ) from awkward._nplikes.numpy_like import NumpyMetadata -from awkward._regularize import is_integer +from awkward._regularize import is_integer, regularize_axis __all__ = ("min", "nanmin") @@ -148,27 +149,29 @@ def nanmin( def _impl(array, axis, keepdims, initial, mask_identity, highlevel, behavior, attrs): + with HighLevelContext(behavior=behavior, attrs=attrs) as ctx: + layout = ctx.unwrap(array, allow_record=False, primitive_policy="error") + out_named_axis = None - if _supports_named_axis(array) and not is_integer(axis): + if _supports_named_axis(ctx) and _check_valid_axis(axis): # Handle named axis # Step 1: Normalize named axis to positional axis axis = _one_axis_to_positional_axis( axis, array.named_axis, array.positional_axis ) - # Step 2: propagate named axis from input to output, - # use strategy "remove one" (see: awkward._namedaxis) - out_named_axis = _identity_named_axis(array.named_axis) - if not keepdims: - out_named_axis = _remove_named_axis(axis, out_named_axis) + # Step 2: propagate named axis from input to output, + # keepdims=True: use strategy "keep all" (see: awkward._namedaxis) + # keepdims=False: use strategy "remove one" (see: awkward._namedaxis) + out_named_axis = _identity_named_axis(array.named_axis) + if not keepdims: + out_named_axis = _remove_named_axis(axis, out_named_axis) - if not isinstance(axis, int): - raise TypeError(f"'axis' must be an integer by now, not {axis!r}") + axis = regularize_axis(axis) - # axis = regularize_axis(axis) # <- is this really needed? + if not is_integer(axis) and axis is not None: + raise TypeError(f"'axis' must be an integer or None by now, not {axis!r}") - with HighLevelContext(behavior=behavior, attrs=attrs) as ctx: - layout = ctx.unwrap(array, allow_record=False, primitive_policy="error") reducer = ak._reducers.Min(initial) out = ak._do.reduce( @@ -179,13 +182,24 @@ def _impl(array, axis, keepdims, initial, mask_identity, highlevel, behavior, at keepdims=keepdims, behavior=ctx.behavior, ) - return ctx.wrap( + + wrapped_out = ctx.wrap( out, highlevel=highlevel, allow_other=True, - named_axis=out_named_axis, ) + if out_named_axis: + # propagate named axis to output + return ak.operations.ak_with_named_axis._impl( + wrapped_out, + named_axis=out_named_axis, + highlevel=highlevel, + behavior=ctx.behavior, + attrs=ctx.attrs, + ) + return wrapped_out + @ak._connect.numpy.implements("amin") @ak._connect.numpy.implements("min") diff --git a/src/awkward/operations/ak_moment.py b/src/awkward/operations/ak_moment.py index 935dad36ba..be6b4958db 100644 --- a/src/awkward/operations/ak_moment.py +++ b/src/awkward/operations/ak_moment.py @@ -11,13 +11,9 @@ ) from awkward._namedaxis import ( AxisName, - _identity_named_axis, - _one_axis_to_positional_axis, - _remove_named_axis, - _supports_named_axis, ) from awkward._nplikes.numpy_like import NumpyMetadata -from awkward._regularize import is_integer +from awkward._regularize import regularize_axis from awkward._typing import Mapping __all__ = ("moment",) @@ -105,32 +101,7 @@ def _impl( behavior: Mapping | None, attrs: Mapping | None, ): - out_named_axis = None - if _supports_named_axis(x) and not is_integer(axis): - # Handle named axis - # Step 1: Convert named axis to positional axis - xaxis = _one_axis_to_positional_axis(axis, x.named_axis, x.positional_axis) - waxis = xaxis - if weight is not None and _supports_named_axis(weight): - waxis = _one_axis_to_positional_axis( - axis, weight.named_axis, weight.positional_axis - ) - if xaxis != waxis: - raise ValueError( - f"ak.mean require the same axis for x and weight, got {xaxis} and {waxis}" - ) - axis = xaxis - - # Step 2: Propagate named axis from input to output using "remove one" strategy - out_named_axis = _identity_named_axis(x.named_axis) - if not keepdims: - # Remove the axis that we are reducing - out_named_axis = _remove_named_axis(axis, out_named_axis) - - if not isinstance(axis, int): - raise TypeError( - f"'axis' must be an integer (or a named axis that maps to an integer) by now, not {axis!r}" - ) + axis = regularize_axis(axis) with HighLevelContext(behavior=behavior, attrs=attrs) as ctx: x_layout, weight_layout = ensure_same_backend( @@ -190,5 +161,4 @@ def _impl( maybe_highlevel_to_lowlevel(sumwxn / sumw), highlevel=highlevel, allow_other=True, - named_axis=out_named_axis, ) diff --git a/src/awkward/operations/ak_nan_to_none.py b/src/awkward/operations/ak_nan_to_none.py index 39529e0fc2..67b42ccda3 100644 --- a/src/awkward/operations/ak_nan_to_none.py +++ b/src/awkward/operations/ak_nan_to_none.py @@ -5,7 +5,6 @@ import awkward as ak from awkward._dispatch import high_level_function from awkward._layout import HighLevelContext -from awkward._namedaxis import _identity_named_axis from awkward._nplikes.numpy_like import NumpyMetadata from awkward._typing import Mapping @@ -64,10 +63,4 @@ def action(layout, continuation, backend, **kwargs): with HighLevelContext(behavior=behavior, attrs=attrs) as ctx: layout = ctx.unwrap(array, allow_record=False, primitive_policy="error") out = ak._do.recursively_apply(layout, action) - return ctx.wrap( - out, - highlevel=highlevel, - named_axis=_identity_named_axis( - array.named_axis - ), # strategy: "keep all" (see: awkward._namedaxis) - ) + return ctx.wrap(out, highlevel=highlevel) diff --git a/src/awkward/operations/ak_nan_to_num.py b/src/awkward/operations/ak_nan_to_num.py index 25105c0ead..fdc7c4b694 100644 --- a/src/awkward/operations/ak_nan_to_num.py +++ b/src/awkward/operations/ak_nan_to_num.py @@ -5,7 +5,6 @@ import awkward as ak from awkward._dispatch import high_level_function from awkward._layout import HighLevelContext, ensure_same_backend -from awkward._namedaxis import _identity_named_axis from awkward._nplikes.numpy_like import NumpyMetadata from awkward._typing import Mapping @@ -153,13 +152,7 @@ def action(inputs, backend, **kwargs): assert isinstance(out, tuple) and len(out) == 1 out = out[0] - return ctx.wrap( - out, - highlevel=highlevel, - named_axis=_identity_named_axis( - array.named_axis - ), # strategy: "keep all" (see: awkward._namedaxis) - ) + return ctx.wrap(out, highlevel=highlevel) @ak._connect.numpy.implements("nan_to_num") diff --git a/src/awkward/operations/ak_num.py b/src/awkward/operations/ak_num.py index a648ea8cdc..4f1eaefc78 100644 --- a/src/awkward/operations/ak_num.py +++ b/src/awkward/operations/ak_num.py @@ -7,13 +7,9 @@ from awkward._layout import HighLevelContext, maybe_posaxis from awkward._namedaxis import ( AxisName, - _identity_named_axis, - _one_axis_to_positional_axis, - _remove_named_axis, - _supports_named_axis, ) from awkward._nplikes.numpy_like import NumpyMetadata -from awkward._regularize import is_integer +from awkward._regularize import is_integer, regularize_axis from awkward._typing import Mapping from awkward.errors import AxisError @@ -105,24 +101,7 @@ def _impl( behavior: Mapping | None, attrs: Mapping | None, ): - out_named_axis = None - if _supports_named_axis(array) and not is_integer(axis): - # Handle named axis - # Step 1: Normalize named axis to positional axis - axis = _one_axis_to_positional_axis( - axis, array.named_axis, array.positional_axis - ) - - # Step 2: propagate named axis from input to output, - # use strategy "remove one" (see: awkward._namedaxis) - out_named_axis = _remove_named_axis( - axis, _identity_named_axis(array.named_axis) - ) - - if not isinstance(axis, int): - raise TypeError(f"'axis' must be an integer by now, not {axis!r}") - - # axis = regularize_axis(axis) # <- is this really needed? + axis = regularize_axis(axis) with HighLevelContext(behavior=behavior, attrs=attrs) as ctx: layout = ctx.unwrap(array, allow_record=False, primitive_policy="error") @@ -148,4 +127,4 @@ def action(layout, depth, **kwargs): out = ak._do.recursively_apply(layout, action, numpy_to_regular=True) - return ctx.wrap(out, highlevel=highlevel, named_axis=out_named_axis) + return ctx.wrap(out, highlevel=highlevel) diff --git a/src/awkward/operations/ak_pad_none.py b/src/awkward/operations/ak_pad_none.py index e15d1b6442..c4f5a73684 100644 --- a/src/awkward/operations/ak_pad_none.py +++ b/src/awkward/operations/ak_pad_none.py @@ -5,9 +5,8 @@ import awkward as ak from awkward._dispatch import high_level_function from awkward._layout import HighLevelContext -from awkward._namedaxis import _supports_named_axis from awkward._nplikes.numpy_like import NumpyMetadata -from awkward._regularize import is_integer, regularize_axis +from awkward._regularize import regularize_axis __all__ = ("pad_none",) @@ -114,12 +113,8 @@ def pad_none( def _impl(array, target, axis, clip, highlevel, behavior, attrs): - out_named_axis = None - if _supports_named_axis(array) and not is_integer(axis): - # Named axis handling - raise NotImplementedError() - axis = regularize_axis(axis) + with HighLevelContext(behavior=behavior, attrs=attrs) as ctx: layout = ctx.unwrap(array, allow_record=False, primitive_policy="error") out = ak._do.pad_none(layout, target, axis, clip=clip) diff --git a/src/awkward/operations/ak_prod.py b/src/awkward/operations/ak_prod.py index 1d7f6c85ff..d2c8623396 100644 --- a/src/awkward/operations/ak_prod.py +++ b/src/awkward/operations/ak_prod.py @@ -6,9 +6,8 @@ from awkward._connect.numpy import UNSUPPORTED from awkward._dispatch import high_level_function from awkward._layout import HighLevelContext -from awkward._namedaxis import _supports_named_axis from awkward._nplikes.numpy_like import NumpyMetadata -from awkward._regularize import is_integer, regularize_axis +from awkward._regularize import regularize_axis __all__ = ("prod", "nanprod") @@ -120,12 +119,8 @@ def nanprod( def _impl(array, axis, keepdims, mask_identity, highlevel, behavior, attrs): - out_named_axis = None - if _supports_named_axis(array) and not is_integer(axis): - # Named axis handling - raise NotImplementedError() - axis = regularize_axis(axis) + with HighLevelContext(behavior=behavior, attrs=attrs) as ctx: layout = ctx.unwrap(array, allow_record=False, primitive_policy="error") reducer = ak._reducers.Prod() diff --git a/src/awkward/operations/ak_ptp.py b/src/awkward/operations/ak_ptp.py index 77a52f4a90..7e69a73bd7 100644 --- a/src/awkward/operations/ak_ptp.py +++ b/src/awkward/operations/ak_ptp.py @@ -10,9 +10,8 @@ maybe_highlevel_to_lowlevel, maybe_posaxis, ) -from awkward._namedaxis import _supports_named_axis from awkward._nplikes.numpy_like import NumpyMetadata -from awkward._regularize import is_integer, regularize_axis +from awkward._regularize import regularize_axis __all__ = ("ptp",) @@ -84,12 +83,8 @@ def ptp( def _impl(array, axis, keepdims, mask_identity, highlevel, behavior, attrs): - out_named_axis = None - if _supports_named_axis(array) and not is_integer(axis): - # Named axis handling - raise NotImplementedError() - axis = regularize_axis(axis) + with HighLevelContext(behavior=behavior, attrs=attrs) as ctx: layout = ctx.unwrap(array, allow_record=False, primitive_policy="error") diff --git a/src/awkward/operations/ak_real.py b/src/awkward/operations/ak_real.py index d42d75a2b7..655e4e8007 100644 --- a/src/awkward/operations/ak_real.py +++ b/src/awkward/operations/ak_real.py @@ -5,7 +5,6 @@ import awkward as ak from awkward._dispatch import high_level_function from awkward._layout import HighLevelContext -from awkward._namedaxis import _supports_named_axis from awkward._nplikes.numpy_like import NumpyMetadata __all__ = ("real",) @@ -38,16 +37,11 @@ def real(array, highlevel=True, behavior=None, attrs=None): def _impl(array, highlevel, behavior, attrs): - out_named_axis = None - if _supports_named_axis(array): - # Named axis handling - raise NotImplementedError() - with HighLevelContext(behavior=behavior, attrs=attrs) as ctx: layout = ctx.unwrap(array, allow_record=False, primitive_policy="error") out = ak._do.recursively_apply(layout, _action_real) - return ctx.wrap(out, highlevel=highlevel, named_axis=out_named_axis) + return ctx.wrap(out, highlevel=highlevel) def _action_real(layout, backend, **kwargs): diff --git a/src/awkward/operations/ak_round.py b/src/awkward/operations/ak_round.py index 9d62c4d074..6d6e0e0471 100644 --- a/src/awkward/operations/ak_round.py +++ b/src/awkward/operations/ak_round.py @@ -6,7 +6,6 @@ from awkward._connect.numpy import UNSUPPORTED from awkward._dispatch import high_level_function from awkward._layout import HighLevelContext -from awkward._namedaxis import _supports_named_axis from awkward._nplikes.numpy_like import NumpyMetadata __all__ = ("round",) @@ -52,11 +51,6 @@ def round( def _impl(array, decimals, highlevel, behavior, attrs): - out_named_axis = None - if _supports_named_axis(array): - # Named axis handling - raise NotImplementedError() - with HighLevelContext(behavior=behavior, attrs=attrs) as ctx: layout = ctx.unwrap(array, allow_record=False, primitive_policy="error") @@ -68,4 +62,4 @@ def action(layout, backend, **kwargs): return None out = ak._do.recursively_apply(layout, action) - return ctx.wrap(out, highlevel=highlevel, named_axis=out_named_axis) + return ctx.wrap(out, highlevel=highlevel) diff --git a/src/awkward/operations/ak_run_lengths.py b/src/awkward/operations/ak_run_lengths.py index d8cf746082..673627fdbc 100644 --- a/src/awkward/operations/ak_run_lengths.py +++ b/src/awkward/operations/ak_run_lengths.py @@ -5,7 +5,6 @@ import awkward as ak from awkward._dispatch import high_level_function from awkward._layout import HighLevelContext -from awkward._namedaxis import _supports_named_axis from awkward._nplikes.numpy_like import NumpyMetadata from awkward._nplikes.shape import unknown_length @@ -99,11 +98,6 @@ def run_lengths(array, *, highlevel=True, behavior=None, attrs=None): def _impl(array, highlevel, behavior, attrs): - out_named_axis = None - if _supports_named_axis(array): - # Named axis handling - raise NotImplementedError() - with HighLevelContext(behavior=behavior, attrs=attrs) as ctx: layout = ctx.unwrap(array, allow_record=False, primitive_policy="error") @@ -230,4 +224,4 @@ def action(layout, **kwargs): return None out = ak._do.recursively_apply(layout, action) - return ctx.wrap(out, highlevel=highlevel, named_axis=out_named_axis) + return ctx.wrap(out, highlevel=highlevel) diff --git a/src/awkward/operations/ak_singletons.py b/src/awkward/operations/ak_singletons.py index d28a4eac71..841ee57bb0 100644 --- a/src/awkward/operations/ak_singletons.py +++ b/src/awkward/operations/ak_singletons.py @@ -5,7 +5,6 @@ import awkward as ak from awkward._dispatch import high_level_function from awkward._layout import HighLevelContext, maybe_posaxis -from awkward._namedaxis import _supports_named_axis from awkward._nplikes.numpy_like import NumpyMetadata from awkward._regularize import is_integer, regularize_axis from awkward.errors import AxisError @@ -57,17 +56,13 @@ def singletons(array, axis=0, *, highlevel=True, behavior=None, attrs=None): def _impl(array, axis, highlevel, behavior, attrs): - out_named_axis = None - if _supports_named_axis(array) and not is_integer(axis): - # Named axis handling - raise NotImplementedError() - axis = regularize_axis(axis) + with HighLevelContext(behavior=behavior, attrs=attrs) as ctx: layout = ctx.unwrap(array, allow_record=False, primitive_policy="error") if not is_integer(axis): - raise TypeError(f"'axis' must be an integer, not {axis!r}") + raise TypeError(f"'axis' must be an integer by now, not {axis!r}") def action(layout, depth, backend, **kwargs): posaxis = maybe_posaxis(layout, axis, depth) diff --git a/src/awkward/operations/ak_softmax.py b/src/awkward/operations/ak_softmax.py index fe83a44cb6..d928566123 100644 --- a/src/awkward/operations/ak_softmax.py +++ b/src/awkward/operations/ak_softmax.py @@ -9,10 +9,9 @@ maybe_highlevel_to_lowlevel, maybe_posaxis, ) -from awkward._namedaxis import _supports_named_axis from awkward._nplikes import ufuncs from awkward._nplikes.numpy_like import NumpyMetadata -from awkward._regularize import is_integer, regularize_axis +from awkward._regularize import regularize_axis __all__ = ("softmax",) @@ -77,11 +76,6 @@ def softmax( def _impl(x, axis, keepdims, mask_identity, highlevel, behavior, attrs): original_axis = axis - out_named_axis = None - if _supports_named_axis(x) and not is_integer(axis): - # Named axis handling - raise NotImplementedError() - axis = regularize_axis(axis) with HighLevelContext(behavior=behavior, attrs=attrs) as ctx: @@ -108,5 +102,4 @@ def _impl(x, axis, keepdims, mask_identity, highlevel, behavior, attrs): maybe_highlevel_to_lowlevel(expx / denom), highlevel=highlevel, allow_other=True, - named_axis=out_named_axis, ) diff --git a/src/awkward/operations/ak_sort.py b/src/awkward/operations/ak_sort.py index d3d83d8e62..c6e38a0648 100644 --- a/src/awkward/operations/ak_sort.py +++ b/src/awkward/operations/ak_sort.py @@ -6,9 +6,8 @@ from awkward._connect.numpy import UNSUPPORTED from awkward._dispatch import high_level_function from awkward._layout import HighLevelContext -from awkward._namedaxis import _supports_named_axis from awkward._nplikes.numpy_like import NumpyMetadata -from awkward._regularize import is_integer, regularize_axis +from awkward._regularize import regularize_axis __all__ = ("sort",) @@ -60,16 +59,12 @@ def sort( def _impl(array, axis, ascending, stable, highlevel, behavior, attrs): - out_named_axis = None - if _supports_named_axis(array) and not is_integer(axis): - # Named axis handling - raise NotImplementedError() - axis = regularize_axis(axis) + with HighLevelContext(behavior=behavior, attrs=attrs) as ctx: layout = ctx.unwrap(array, allow_record=False, primitive_policy="error") out = ak._do.sort(layout, axis, ascending, stable) - return ctx.wrap(out, highlevel=highlevel, named_axis=out_named_axis) + return ctx.wrap(out, highlevel=highlevel) @ak._connect.numpy.implements("sort") diff --git a/src/awkward/operations/ak_std.py b/src/awkward/operations/ak_std.py index 1d6db0a9c6..b6972067dc 100644 --- a/src/awkward/operations/ak_std.py +++ b/src/awkward/operations/ak_std.py @@ -11,10 +11,9 @@ maybe_highlevel_to_lowlevel, maybe_posaxis, ) -from awkward._namedaxis import _supports_named_axis from awkward._nplikes import ufuncs from awkward._nplikes.numpy_like import NumpyMetadata -from awkward._regularize import is_integer, regularize_axis +from awkward._regularize import regularize_axis __all__ = ("std", "nanstd") @@ -166,11 +165,6 @@ def nanstd( def _impl(x, weight, ddof, axis, keepdims, mask_identity, highlevel, behavior, attrs): - out_named_axis = None - if _supports_named_axis(x) and not is_integer(axis): - # Named axis handling - raise NotImplementedError() - axis = regularize_axis(axis) with HighLevelContext(behavior=behavior, attrs=attrs) as ctx: @@ -225,7 +219,6 @@ def _impl(x, weight, ddof, axis, keepdims, mask_identity, highlevel, behavior, a maybe_highlevel_to_lowlevel(out), highlevel=highlevel, allow_other=True, - named_axis=out_named_axis, ) diff --git a/src/awkward/operations/ak_strings_astype.py b/src/awkward/operations/ak_strings_astype.py index cd864ce67a..b0834db3a6 100644 --- a/src/awkward/operations/ak_strings_astype.py +++ b/src/awkward/operations/ak_strings_astype.py @@ -5,7 +5,6 @@ import awkward as ak from awkward._dispatch import high_level_function from awkward._layout import HighLevelContext -from awkward._namedaxis import _supports_named_axis from awkward._nplikes.numpy import Numpy from awkward._nplikes.numpy_like import NumpyMetadata @@ -59,11 +58,6 @@ def strings_astype(array, to, *, highlevel=True, behavior=None, attrs=None): def _impl(array, to, highlevel, behavior, attrs): - out_named_axis = None - if _supports_named_axis(array): - # Named axis handling - raise NotImplementedError() - def action(layout, **kwargs): if layout.is_list and ( layout.parameter("__array__") == "string" @@ -89,4 +83,4 @@ def action(layout, **kwargs): with HighLevelContext(behavior=behavior, attrs=attrs) as ctx: layout = ctx.unwrap(array, allow_record=False, primitive_policy="error") out = ak._do.recursively_apply(layout, action) - return ctx.wrap(out, highlevel=highlevel, named_axis=out_named_axis) + return ctx.wrap(out, highlevel=highlevel) diff --git a/src/awkward/operations/ak_sum.py b/src/awkward/operations/ak_sum.py index a7b969eca8..add5bc6f9c 100644 --- a/src/awkward/operations/ak_sum.py +++ b/src/awkward/operations/ak_sum.py @@ -6,7 +6,13 @@ from awkward._connect.numpy import UNSUPPORTED from awkward._dispatch import high_level_function from awkward._layout import HighLevelContext -from awkward._namedaxis import _supports_named_axis +from awkward._namedaxis import ( + _check_valid_axis, + _identity_named_axis, + _one_axis_to_positional_axis, + _remove_named_axis, + _supports_named_axis, +) from awkward._nplikes.numpy_like import NumpyMetadata from awkward._regularize import is_integer, regularize_axis @@ -270,14 +276,29 @@ def nansum( def _impl(array, axis, keepdims, mask_identity, highlevel, behavior, attrs): + with HighLevelContext(behavior=behavior, attrs=attrs) as ctx: + layout = ctx.unwrap(array, allow_record=False, primitive_policy="error") + out_named_axis = None - if _supports_named_axis(array) and not is_integer(axis): - # Named axis handling - raise NotImplementedError() + if _supports_named_axis(ctx) and _check_valid_axis(axis): + # Handle named axis + # Step 1: Normalize named axis to positional axis + axis = _one_axis_to_positional_axis( + axis, array.named_axis, array.positional_axis + ) + + # Step 2: propagate named axis from input to output, + # keepdims=True: use strategy "keep all" (see: awkward._namedaxis) + # keepdims=False: use strategy "remove one" (see: awkward._namedaxis) + out_named_axis = _identity_named_axis(array.named_axis) + if not keepdims: + out_named_axis = _remove_named_axis(axis, out_named_axis) axis = regularize_axis(axis) - with HighLevelContext(behavior=behavior, attrs=attrs) as ctx: - layout = ctx.unwrap(array, allow_record=False, primitive_policy="error") + + if not is_integer(axis) and axis is not None: + raise TypeError(f"'axis' must be an integer or None by now, not {axis!r}") + reducer = ak._reducers.Sum() out = ak._do.reduce( @@ -288,13 +309,24 @@ def _impl(array, axis, keepdims, mask_identity, highlevel, behavior, attrs): keepdims=keepdims, behavior=ctx.behavior, ) - return ctx.wrap( + + wrapped_out = ctx.wrap( out, highlevel=highlevel, allow_other=True, - named_axis=out_named_axis, ) + if out_named_axis: + # propagate named axis to output + return ak.operations.ak_with_named_axis._impl( + wrapped_out, + named_axis=out_named_axis, + highlevel=highlevel, + behavior=ctx.behavior, + attrs=ctx.attrs, + ) + return wrapped_out + @ak._connect.numpy.implements("sum") def _nep_18_impl_sum( diff --git a/src/awkward/operations/ak_transform.py b/src/awkward/operations/ak_transform.py index 8d3425fb49..23b4dbfd4e 100644 --- a/src/awkward/operations/ak_transform.py +++ b/src/awkward/operations/ak_transform.py @@ -15,7 +15,6 @@ ) from awkward._dispatch import high_level_function from awkward._layout import HighLevelContext, ensure_same_backend -from awkward._namedaxis import _supports_named_axis __all__ = ("transform",) @@ -469,11 +468,6 @@ def _impl( highlevel, attrs, ): - out_named_axis = None - if _supports_named_axis(array): - # Named axis handling - raise NotImplementedError() - with HighLevelContext(behavior=behavior, attrs=attrs) as ctx: layouts = ensure_same_backend( ctx.unwrap( @@ -608,8 +602,6 @@ def action(inputs, **kwargs): "or tuple of Contents, but instead only returned None." ) elif len(out) == 1: - return ctx.wrap(out[0], highlevel=highlevel, named_axis=out_named_axis) + return ctx.wrap(out[0], highlevel=highlevel) else: - return tuple( - ctx.wrap(x, highlevel=highlevel, named_axis=out_named_axis) for x in out - ) + return tuple(ctx.wrap(x, highlevel=highlevel) for x in out) diff --git a/src/awkward/operations/ak_unflatten.py b/src/awkward/operations/ak_unflatten.py index 053ef69365..a4e659a197 100644 --- a/src/awkward/operations/ak_unflatten.py +++ b/src/awkward/operations/ak_unflatten.py @@ -6,7 +6,11 @@ from awkward._backends.numpy import NumpyBackend from awkward._dispatch import high_level_function from awkward._layout import HighLevelContext, ensure_same_backend, maybe_posaxis -from awkward._namedaxis import _supports_named_axis +from awkward._namedaxis import ( + _check_valid_axis, + _one_axis_to_positional_axis, + _supports_named_axis, +) from awkward._nplikes.numpy_like import NumpyMetadata from awkward._nplikes.shape import unknown_length from awkward._nplikes.typetracer import is_unknown_scalar @@ -92,13 +96,6 @@ def unflatten(array, counts, axis=0, *, highlevel=True, behavior=None, attrs=Non def _impl(array, counts, axis, highlevel, behavior, attrs): - out_named_axis = None - if _supports_named_axis(array) and not is_integer(axis): - # Named axis handling - raise NotImplementedError() - - axis = regularize_axis(axis) - with HighLevelContext(behavior=behavior, attrs=attrs) as ctx: layout, maybe_counts_layout = ensure_same_backend( ctx.unwrap(array, allow_record=False, primitive_policy="error"), @@ -111,6 +108,22 @@ def _impl(array, counts, axis, highlevel, behavior, attrs): ), ) + if _supports_named_axis(ctx) and _check_valid_axis(axis): + # Handle named axis + # Step 1: Normalize named axis to positional axis + axis = _one_axis_to_positional_axis( + axis, array.named_axis, array.positional_axis + ) + + # Step 2: propagate named axis from input to output, + # use strategy "remove all" (see: awkward._namedaxis) + out_named_axis = None + + axis = regularize_axis(axis) + + if not is_integer(axis): + raise TypeError(f"'axis' must be an integer by now, not {axis!r}") + if is_integer_like(maybe_counts_layout): # Regularize unknown values to unknown lengths if ( @@ -298,4 +311,16 @@ def apply(layout, depth, backend, **kwargs): f"at axis={axis}" ) - return ctx.wrap(out, highlevel=highlevel, named_axis=out_named_axis) + wrapped_out = ctx.wrap( + out, + highlevel=highlevel, + ) + + # propagate named axis to output + return ak.operations.ak_with_named_axis._impl( + wrapped_out, + named_axis=out_named_axis, + highlevel=highlevel, + behavior=ctx.behavior, + attrs=ctx.attrs, + ) diff --git a/src/awkward/operations/ak_unzip.py b/src/awkward/operations/ak_unzip.py index 86c3a11437..8c19380133 100644 --- a/src/awkward/operations/ak_unzip.py +++ b/src/awkward/operations/ak_unzip.py @@ -5,7 +5,6 @@ import awkward as ak from awkward._dispatch import high_level_function from awkward._layout import HighLevelContext -from awkward._namedaxis import _supports_named_axis from awkward._nplikes.numpy_like import NumpyMetadata __all__ = ("unzip",) @@ -50,13 +49,9 @@ def unzip(array, *, highlevel=True, behavior=None, attrs=None): def _impl(array, highlevel, behavior, attrs): - out_named_axis = None - if _supports_named_axis(array): - # Named axis handling - raise NotImplementedError() - with HighLevelContext(behavior=behavior, attrs=attrs) as ctx: layout = ctx.unwrap(array, allow_record=True, primitive_policy="error") + fields = ak.operations.fields(layout) def check_for_union(layout, **kwargs): @@ -73,18 +68,13 @@ def check_for_union(layout, **kwargs): ak._do.recursively_apply(layout, check_for_union, return_array=False) if len(fields) == 0: - return ( - ctx.wrap( - layout, highlevel=highlevel, allow_other=True, named_axis=out_named_axis - ), - ) + return (ctx.wrap(layout, highlevel=highlevel, allow_other=True),) else: return tuple( ctx.wrap( layout[n], highlevel=highlevel, allow_other=True, - named_axis=out_named_axis, ) for n in fields ) diff --git a/src/awkward/operations/ak_values_astype.py b/src/awkward/operations/ak_values_astype.py index dcfb97e67f..fa25ca5a35 100644 --- a/src/awkward/operations/ak_values_astype.py +++ b/src/awkward/operations/ak_values_astype.py @@ -5,7 +5,6 @@ import awkward as ak from awkward._dispatch import high_level_function from awkward._layout import HighLevelContext -from awkward._namedaxis import _supports_named_axis from awkward._nplikes.numpy_like import NumpyMetadata __all__ = ("values_astype",) @@ -71,13 +70,9 @@ def values_astype( def _impl(array, to, including_unknown, highlevel, behavior, attrs): - out_named_axis = None - if _supports_named_axis(array): - # Named axis handling - raise NotImplementedError() - with HighLevelContext(behavior=behavior, attrs=attrs) as ctx: layout = ctx.unwrap(array, allow_record=False, primitive_policy="error") + to_str = ak.types.numpytype.dtype_to_primitive(np.dtype(to)) out = ak._do.numbers_to_type(layout, to_str, including_unknown) - return ctx.wrap(out, highlevel=highlevel, named_axis=out_named_axis) + return ctx.wrap(out, highlevel=highlevel) diff --git a/src/awkward/operations/ak_var.py b/src/awkward/operations/ak_var.py index eaedd21c4a..521ab9ccad 100644 --- a/src/awkward/operations/ak_var.py +++ b/src/awkward/operations/ak_var.py @@ -11,9 +11,8 @@ maybe_highlevel_to_lowlevel, maybe_posaxis, ) -from awkward._namedaxis import _supports_named_axis from awkward._nplikes.numpy_like import NumpyMetadata -from awkward._regularize import is_integer, regularize_axis +from awkward._regularize import regularize_axis __all__ = ("var", "nanvar") @@ -171,11 +170,6 @@ def nanvar( def _impl(x, weight, ddof, axis, keepdims, mask_identity, highlevel, behavior, attrs): - out_named_axis = None - if _supports_named_axis(x) and not is_integer(axis): - # Named axis handling - raise NotImplementedError() - axis = regularize_axis(axis) with HighLevelContext(behavior=behavior, attrs=attrs) as ctx: @@ -277,7 +271,6 @@ def _impl(x, weight, ddof, axis, keepdims, mask_identity, highlevel, behavior, a maybe_highlevel_to_lowlevel(out), highlevel=highlevel, allow_other=True, - named_axis=out_named_axis, ) diff --git a/src/awkward/operations/ak_with_field.py b/src/awkward/operations/ak_with_field.py index a3a18e0ab1..671a061978 100644 --- a/src/awkward/operations/ak_with_field.py +++ b/src/awkward/operations/ak_with_field.py @@ -7,7 +7,6 @@ import awkward as ak from awkward._dispatch import high_level_function from awkward._layout import HighLevelContext, ensure_same_backend -from awkward._namedaxis import _supports_named_axis from awkward._nplikes.numpy_like import NumpyMetadata from awkward._regularize import is_non_string_like_sequence @@ -50,11 +49,6 @@ def with_field(array, what, where=None, *, highlevel=True, behavior=None, attrs= def _impl(base, what, where, highlevel, behavior, attrs): - out_named_axis = None - if _supports_named_axis(base): - # Named axis handling - raise NotImplementedError() - if not ( where is None or isinstance(where, str) @@ -167,4 +161,4 @@ def action(inputs, **kwargs): assert isinstance(out, tuple) and len(out) == 1 - return ctx.wrap(out[0], highlevel=highlevel, named_axis=out_named_axis) + return ctx.wrap(out[0], highlevel=highlevel) diff --git a/src/awkward/operations/ak_with_name.py b/src/awkward/operations/ak_with_name.py index e908f9581e..34a9e1cb89 100644 --- a/src/awkward/operations/ak_with_name.py +++ b/src/awkward/operations/ak_with_name.py @@ -5,7 +5,6 @@ import awkward as ak from awkward._dispatch import high_level_function from awkward._layout import HighLevelContext -from awkward._namedaxis import _supports_named_axis from awkward._nplikes.numpy_like import NumpyMetadata __all__ = ("with_name",) @@ -47,11 +46,6 @@ def with_name(array, name, *, highlevel=True, behavior=None, attrs=None): def _impl(array, name, highlevel, behavior, attrs): - out_named_axis = None - if _supports_named_axis(array): - # Named axis handling - raise NotImplementedError() - with HighLevelContext(behavior=behavior, attrs=attrs) as ctx: layout = ctx.unwrap(array, allow_record=True, primitive_policy="error") @@ -68,4 +62,4 @@ def action(layout, **ignore): out = ak._do.recursively_apply(layout, action) - return ctx.wrap(out, highlevel=highlevel, named_axis=out_named_axis) + return ctx.wrap(out, highlevel=highlevel) diff --git a/src/awkward/operations/ak_with_named_axis.py b/src/awkward/operations/ak_with_named_axis.py index 1a689b47e8..c6bbd07c95 100644 --- a/src/awkward/operations/ak_with_named_axis.py +++ b/src/awkward/operations/ak_with_named_axis.py @@ -58,14 +58,13 @@ def with_named_axis( def _impl(array, named_axis, highlevel, behavior, attrs): - if not named_axis: # no-op, e.g. if named_axis is None or () or {} - return array - with HighLevelContext(behavior=behavior, attrs=attrs) as ctx: layout = ctx.unwrap(array, allow_record=False) # Named axis handling ndim = layout.purelist_depth + if not named_axis: # no-op, e.g. named_axis is None, (), {} + named_axis = (None,) * ndim if isinstance(named_axis, dict): _named_axis = tuple(named_axis.get(i, None) for i in range(ndim)) for k, i in named_axis.items(): @@ -90,4 +89,4 @@ def _impl(array, named_axis, highlevel, behavior, attrs): ) out_ctx = HighLevelContext(behavior=ctx.behavior, attrs=attrs).finalize() - return out_ctx.wrap(layout, highlevel=highlevel) + return out_ctx.wrap(layout, highlevel=highlevel, allow_other=True) diff --git a/src/awkward/operations/ak_with_parameter.py b/src/awkward/operations/ak_with_parameter.py index 57a0793c76..9349fe3a1f 100644 --- a/src/awkward/operations/ak_with_parameter.py +++ b/src/awkward/operations/ak_with_parameter.py @@ -4,7 +4,6 @@ from awkward._dispatch import high_level_function from awkward._layout import HighLevelContext -from awkward._namedaxis import _supports_named_axis from awkward._nplikes.numpy_like import NumpyMetadata __all__ = ("with_parameter",) @@ -45,13 +44,8 @@ def with_parameter( def _impl(array, parameter, value, highlevel, behavior, attrs): - out_named_axis = None - if _supports_named_axis(array): - # Named axis handling - raise NotImplementedError() - with HighLevelContext(behavior=behavior, attrs=attrs) as ctx: layout = ctx.unwrap(array, allow_record=True, primitive_policy="error") out = layout.with_parameter(parameter, value) - return ctx.wrap(out, highlevel=highlevel, named_axis=out_named_axis) + return ctx.wrap(out, highlevel=highlevel) diff --git a/src/awkward/operations/ak_without_field.py b/src/awkward/operations/ak_without_field.py index 0f0976fba8..2fc06da053 100644 --- a/src/awkward/operations/ak_without_field.py +++ b/src/awkward/operations/ak_without_field.py @@ -7,7 +7,6 @@ import awkward as ak from awkward._dispatch import high_level_function from awkward._layout import HighLevelContext -from awkward._namedaxis import _supports_named_axis from awkward._nplikes.numpy_like import NumpyMetadata __all__ = ("without_field",) @@ -47,11 +46,6 @@ def without_field(array, where, *, highlevel=True, behavior=None, attrs=None): def _impl(base, where, highlevel, behavior, attrs): - out_named_axis = None - if _supports_named_axis(base): - # Named axis handling - raise NotImplementedError() - if isinstance(where, str): where = [where] elif not (isinstance(where, Sequence) and all(isinstance(x, str) for x in where)): @@ -101,4 +95,4 @@ def action(layout, depth_context, **kwargs): return None out = ak._do.recursively_apply(base, action, depth_context={"where": where}) - return ctx.wrap(out, highlevel=highlevel, named_axis=out_named_axis) + return ctx.wrap(out, highlevel=highlevel) diff --git a/src/awkward/operations/ak_without_parameters.py b/src/awkward/operations/ak_without_parameters.py index 772e78a0ac..f414915a45 100644 --- a/src/awkward/operations/ak_without_parameters.py +++ b/src/awkward/operations/ak_without_parameters.py @@ -5,7 +5,6 @@ import awkward as ak from awkward._dispatch import high_level_function from awkward._layout import HighLevelContext -from awkward._namedaxis import _supports_named_axis from awkward._nplikes.numpy_like import NumpyMetadata __all__ = ("without_parameters",) @@ -39,11 +38,6 @@ def without_parameters(array, *, highlevel=True, behavior=None, attrs=None): def _impl(array, highlevel, behavior, attrs): - out_named_axis = None - if _supports_named_axis(array): - # Named axis handling - raise NotImplementedError() - with HighLevelContext(behavior=behavior, attrs=attrs) as ctx: layout = ctx.unwrap(array, allow_record=False, primitive_policy="error") @@ -51,4 +45,4 @@ def _impl(array, highlevel, behavior, attrs): layout, lambda layout, **kwargs: None, keep_parameters=False ) - return ctx.wrap(out, highlevel=highlevel, named_axis=out_named_axis) + return ctx.wrap(out, highlevel=highlevel) diff --git a/src/awkward/operations/ak_zip.py b/src/awkward/operations/ak_zip.py index 8ace5c2b22..02deb422f0 100644 --- a/src/awkward/operations/ak_zip.py +++ b/src/awkward/operations/ak_zip.py @@ -7,7 +7,6 @@ import awkward as ak from awkward._dispatch import high_level_function from awkward._layout import HighLevelContext, ensure_same_backend -from awkward._namedaxis import _supports_named_axis from awkward._nplikes.numpy_like import NumpyMetadata __all__ = ("zip",) @@ -176,11 +175,6 @@ def _impl( if depth_limit is not None and depth_limit <= 0: raise ValueError("depth_limit must be None or at least 1") - out_named_axis = None - if _supports_named_axis(array): - # Named axis handling - raise NotImplementedError() - with HighLevelContext(behavior=behavior, attrs=attrs) as ctx: if isinstance(arrays, Mapping): layouts = ensure_same_backend( @@ -255,4 +249,4 @@ def action(inputs, depth, backend, **ignore): out = out[0] assert isinstance(out, ak.record.Record) - return ctx.wrap(out, highlevel=highlevel, named_axis=out_named_axis) + return ctx.wrap(out, highlevel=highlevel) diff --git a/tests/test_2596_named_axis.py b/tests/test_2596_named_axis.py index 01899ec74a..c226c1328f 100644 --- a/tests/test_2596_named_axis.py +++ b/tests/test_2596_named_axis.py @@ -305,7 +305,16 @@ def test_named_axis_ak_cartesian(): def test_named_axis_ak_categories(): - assert True + array = ak.str.to_categorical([["one", "two"], ["one", "three"], ["one", "four"]]) + + named_array = ak.with_named_axis(array, named_axis=("a", "b")) + + # assert ak.all(ak.categories(array) == ak.categories(named_array)) # FIX: ufuncs + assert ( + ak.categories(array).named_axis + == ak.categories(named_array).named_axis + == (None,) + ) def test_named_axis_ak_combinations(): @@ -325,39 +334,148 @@ def test_named_axis_ak_copy(): def test_named_axis_ak_corr(): - assert True + array_x = ak.Array([[0, 1.1], [3.3, 4.4]]) + array_y = ak.Array([[0, 1], [3, 4]]) + + named_array_x = ak.with_named_axis(array_x, ("x", "y")) + named_array_y = ak.with_named_axis(array_y, ("x", "y")) + + assert ak.all( + ak.corr(array_x, array_y, axis=0) + == ak.corr(named_array_x, named_array_y, axis="x") + ) + assert ak.all( + ak.corr(array_x, array_y, axis=1) + == ak.corr(named_array_x, named_array_y, axis="y") + ) + assert ak.all( + ak.corr(array_x, array_y, axis=None) + == ak.corr(named_array_x, named_array_y, axis=None) + ) + + assert ak.corr(named_array_x, named_array_y, axis="x").named_axis == ("y",) + assert ak.corr(named_array_x, named_array_y, axis="y").named_axis == ("x",) + assert ak.corr(named_array_x, named_array_y, axis=None).named_axis == (None,) def test_named_axis_ak_count(): - assert True + array = ak.Array([[1, 2], [3], [], [4, 5, 6]]) + + named_array = ak.with_named_axis(array, ("x", "y")) + + assert ak.all(ak.count(array, axis=0) == ak.count(named_array, axis="x")) + assert ak.all(ak.count(array, axis=1) == ak.count(named_array, axis="y")) + assert ak.all(ak.count(array, axis=None) == ak.count(named_array, axis=None)) + + assert ak.count(named_array, axis="x").named_axis == ("y",) + assert ak.count(named_array, axis="y").named_axis == ("x",) + assert ak.count(named_array, axis=None).named_axis == (None,) def test_named_axis_ak_count_nonzero(): - assert True + array = ak.Array([[1, 2], [3], [], [4, 5, 6]]) + + named_array = ak.with_named_axis(array, ("x", "y")) + + assert ak.all( + ak.count_nonzero(array, axis=0) == ak.count_nonzero(named_array, axis="x") + ) + assert ak.all( + ak.count_nonzero(array, axis=1) == ak.count_nonzero(named_array, axis="y") + ) + assert ak.all( + ak.count_nonzero(array, axis=None) == ak.count_nonzero(named_array, axis=None) + ) + + assert ak.count_nonzero(named_array, axis="x").named_axis == ("y",) + assert ak.count_nonzero(named_array, axis="y").named_axis == ("x",) + assert ak.count_nonzero(named_array, axis=None).named_axis == (None,) def test_named_axis_ak_covar(): - assert True + array_x = ak.Array([[0, 1.1], [3.3, 4.4]]) + array_y = ak.Array([[0, 1], [3, 4]]) + + named_array_x = ak.with_named_axis(array_x, ("x", "y")) + named_array_y = ak.with_named_axis(array_y, ("x", "y")) + + assert ak.all( + ak.covar(array_x, array_y, axis=0) + == ak.covar(named_array_x, named_array_y, axis="x") + ) + assert ak.all( + ak.covar(array_x, array_y, axis=1) + == ak.covar(named_array_x, named_array_y, axis="y") + ) + assert ak.all( + ak.covar(array_x, array_y, axis=None) + == ak.covar(named_array_x, named_array_y, axis=None) + ) + + assert ak.covar(named_array_x, named_array_y, axis="x").named_axis == ("y",) + assert ak.covar(named_array_x, named_array_y, axis="y").named_axis == ("x",) + assert ak.covar(named_array_x, named_array_y, axis=None).named_axis == (None,) def test_named_axis_ak_drop_none(): - assert True + array = ak.Array([[1, None], [3], [None], [4, None, 6]]) + + named_array = ak.with_named_axis(array, ("x", "y")) + + assert ak.all(ak.drop_none(array, axis=0) == ak.drop_none(named_array, axis="x")) + assert ak.all(ak.drop_none(array, axis=1) == ak.drop_none(named_array, axis="y")) + assert ak.all( + ak.drop_none(array, axis=None) == ak.drop_none(named_array, axis=None) + ) + + assert ak.drop_none(named_array, axis="x").named_axis == ("x", "y") + assert ak.drop_none(named_array, axis="y").named_axis == ("x", "y") + assert ak.drop_none(named_array, axis=None).named_axis == ("x", "y") def test_named_axis_ak_enforce_type(): - assert True + array = ak.Array([[1, 2], [3], [], [4, 5, 6]]) + + named_array = ak.with_named_axis(array, ("x", "y")) + + assert ak.enforce_type(named_array, "var * ?int64").named_axis == ("x", "y") def test_named_axis_ak_fields(): + # skip assert True def test_named_axis_ak_fill_none(): - assert True + array = ak.Array([[1.1, None, 2.2], [], [None, 3.3, 4.4]]) + + named_array = ak.with_named_axis(array, ("x", "y")) + + assert ak.all( + ak.fill_none(array, 0, axis=0) == ak.fill_none(named_array, 0, axis="x") + ) + assert ak.all( + ak.fill_none(array, 0, axis=1) == ak.fill_none(named_array, 0, axis="y") + ) + assert ak.all( + ak.fill_none(array, 0, axis=None) == ak.fill_none(named_array, 0, axis=None) + ) + + assert ak.fill_none(named_array, 0, axis="x").named_axis == ("x", "y") + assert ak.fill_none(named_array, 0, axis="y").named_axis == ("x", "y") + assert ak.fill_none(named_array, 0, axis=None).named_axis == ("x", "y") def test_named_axis_ak_firsts(): - assert True + array = ak.Array([[1.1], [2.2], [], [3.3], [], [], [4.4], [5.5]]) + + named_array = ak.with_named_axis(array, ("x", "y")) + + assert ak.all(ak.firsts(array, axis=0) == ak.firsts(named_array, axis="x")) + assert ak.all(ak.firsts(array, axis=1) == ak.firsts(named_array, axis="y")) + + assert ak.firsts(named_array, axis="x").named_axis == ("x",) + assert ak.firsts(named_array, axis="y").named_axis == ("y",) def test_named_axis_ak_flatten(): @@ -365,103 +483,156 @@ def test_named_axis_ak_flatten(): def test_named_axis_ak_from_arrow(): + # skip assert True def test_named_axis_ak_from_arrow_schema(): + # skip assert True def test_named_axis_ak_from_avro_file(): + # skip assert True def test_named_axis_ak_from_buffers(): + # skip assert True def test_named_axis_ak_from_categorical(): + # skip assert True def test_named_axis_ak_from_cupy(): + # skip assert True def test_named_axis_ak_from_dlpack(): + # skip assert True def test_named_axis_ak_from_feather(): + # skip assert True def test_named_axis_ak_from_iter(): + # skip assert True def test_named_axis_ak_from_jax(): + # skip assert True def test_named_axis_ak_from_json(): + # skip assert True def test_named_axis_ak_from_numpy(): + # skip assert True def test_named_axis_ak_from_parquet(): + # skip assert True def test_named_axis_ak_from_raggedtensor(): + # skip assert True def test_named_axis_ak_from_rdataframe(): + # skip assert True def test_named_axis_ak_from_regular(): + # skip assert True def test_named_axis_ak_full_like(): + # skip assert True def test_named_axis_ak_imag(): - assert True + array = ak.Array([[1 + 2j], [2 + 1j], []]) + + named_array = ak.with_named_axis(array, ("x", "y")) + + assert ak.all(ak.imag(array) == ak.imag(named_array)) + assert ak.imag(named_array).named_axis == ("x", "y") def test_named_axis_ak_is_categorical(): + # skip assert True def test_named_axis_ak_is_none(): - assert True + array = ak.Array([[1, None], [3], [None], [4, None, 6]]) + + named_array = ak.with_named_axis(array, ("x", "y")) + + assert ak.all(ak.is_none(array, axis=0) == ak.is_none(named_array, axis="x")) + assert ak.all(ak.is_none(array, axis=1) == ak.is_none(named_array, axis="y")) + + assert ak.is_none(named_array, axis="x").named_axis == ("x", "y") + assert ak.is_none(named_array, axis="y").named_axis == ("x", "y") def test_named_axis_ak_is_tuple(): + # skip assert True def test_named_axis_ak_is_valid(): + # skip assert True def test_named_axis_ak_isclose(): + # skip assert True def test_named_axis_ak_linear_fit(): + # skip assert True def test_named_axis_ak_local_index(): - assert True + array = ak.Array( + [[[0.0, 1.1, 2.2], []], [[3.3, 4.4]], [], [[5.5], [], [6.6, 7.7, 8.8, 9.9]]] + ) + + named_array = ak.with_named_axis(array, ("x", "y", "z")) + + assert ak.all( + ak.local_index(array, axis=0) == ak.local_index(named_array, axis="x") + ) + assert ak.all( + ak.local_index(array, axis=1) == ak.local_index(named_array, axis="y") + ) + assert ak.all( + ak.local_index(array, axis=2) == ak.local_index(named_array, axis="z") + ) + + assert ak.local_index(named_array, axis="x").named_axis == ("x",) + assert ak.local_index(named_array, axis="y").named_axis == ("x", "y") + assert ak.local_index(named_array, axis="z").named_axis == ("x", "y", "z") def test_named_axis_ak_mask(): @@ -492,23 +663,55 @@ def test_named_axis_ak_max(): def test_named_axis_ak_mean(): - assert True + array = ak.Array([[1, 2], [3], [4, 5, 6]]) + + named_array = ak.with_named_axis(array, ("x", "y")) + + assert ak.all(ak.mean(array, axis=0) == ak.mean(named_array, axis="x")) + assert ak.all(ak.mean(array, axis=1) == ak.mean(named_array, axis="y")) + assert ak.mean(array, axis=None) == ak.mean(named_array, axis=None) + + assert ak.mean(named_array, axis="x").named_axis == ("y",) + assert ak.mean(named_array, axis="y").named_axis == ("x",) + assert ak.mean(named_array, axis=None).named_axis == (None,) def test_named_axis_ak_merge_option_of_records(): + # skip assert True def test_named_axis_ak_merge_union_of_records(): + # skip assert True def test_named_axis_ak_metadata_from_parquet(): + # skip assert True def test_named_axis_ak_min(): - assert True + array = ak.Array([[1, 2], [3], [], [4, 5, 6]]) + + named_array = ak.with_named_axis(array, named_axis=("events", "jets")) + + # first check that they work the same + assert ak.all(ak.min(array, axis=0) == ak.min(named_array, axis="events")) + assert ak.all(ak.min(array, axis=1) == ak.min(named_array, axis="jets")) + + # check that result axis names are correctly propagated + assert ( + ak.min(named_array, axis=0).named_axis + == ak.min(named_array, axis="events").named_axis + == ("jets",) + ) + assert ( + ak.min(named_array, axis=1).named_axis + == ak.min(named_array, axis="jets").named_axis + == ("events",) + ) + assert ak.min(named_array, axis=None).named_axis == (None,) def test_named_axis_ak_moment(): @@ -584,14 +787,26 @@ def test_named_axis_ak_strings_astype(): def test_named_axis_ak_sum(): - assert True + array = ak.Array([[1, 2], [3], [], [4, 5, 6]]) + + named_array = ak.with_named_axis(array, ("x", "y")) + + assert ak.all(ak.sum(array, axis=0) == ak.sum(named_array, axis="x")) + assert ak.all(ak.sum(array, axis=1) == ak.sum(named_array, axis="y")) + assert ak.sum(array, axis=None) == ak.sum(named_array, axis=None) + + assert ak.sum(named_array, axis="x").named_axis == ("y",) + assert ak.sum(named_array, axis="y").named_axis == ("x",) + assert ak.sum(named_array, axis=None).named_axis == (None,) def test_named_axis_ak_to_arrow(): + # skip assert True def test_named_axis_ak_to_arrow_table(): + # skip assert True @@ -600,38 +815,47 @@ def test_named_axis_ak_to_backend(): def test_named_axis_ak_to_buffers(): + # skip assert True def test_named_axis_ak_to_cupy(): + # skip assert True def test_named_axis_ak_to_dataframe(): + # skip assert True def test_named_axis_ak_to_feather(): + # skip assert True def test_named_axis_ak_to_jax(): + # skip assert True def test_named_axis_ak_to_json(): + # skip assert True def test_named_axis_ak_to_layout(): + # skip assert True def test_named_axis_ak_to_list(): + # skip assert True def test_named_axis_ak_to_numpy(): + # skip assert True @@ -640,46 +864,66 @@ def test_named_axis_ak_to_packed(): def test_named_axis_ak_to_parquet(): + # skip assert True def test_named_axis_ak_to_parquet_dataset(): + # skip assert True def test_named_axis_ak_to_parquet_row_groups(): + # skip assert True def test_named_axis_ak_to_raggedtensor(): + # skip assert True def test_named_axis_ak_to_rdataframe(): + # skip assert True def test_named_axis_ak_to_regular(): + # skip assert True def test_named_axis_ak_transform(): + # skip assert True def test_named_axis_ak_type(): + # skip assert True def test_named_axis_ak_unflatten(): - assert True + array = ak.Array([[1, 2, 3, 4], [], [5, 6, 7], [8, 9]]) + + named_array = ak.with_named_axis(array, ("x", "y")) + + counts = ak.Array([2, 2, 1, 2, 1, 1]) + + assert ak.all( + ak.unflatten(array, counts, axis=1) + == ak.unflatten(named_array, counts, axis="y") + ) + assert ak.unflatten(named_array, counts, axis="y").named_axis == (None, None, None) def test_named_axis_ak_unzip(): + # skip assert True def test_named_axis_ak_validity_error(): + # skip assert True