docs: document named-axis support for axis in high-level operations by aashirvad08 · Pull Request #3823 · scikit-hep/awkward

aashirvad08 · 2026-01-25T08:32:36Z

This PR improves the docstrings of all highlevel functions to state that they support named axes.

Aligns the docstring with named-axis support
No runtime behavior changes

codecov · 2026-01-25T08:37:45Z

Codecov Report

✅ All modified and coverable lines are covered by tests.
✅ Project coverage is 82.65%. Comparing base (02f3080) to head (a83d900).
⚠️ Report is 8 commits behind head on main.

Additional details and impacted files

Files with missing lines	Coverage Δ
src/awkward/operations/ak_all.py	`96.66% <ø> (ø)`
src/awkward/operations/ak_any.py	`96.66% <ø> (ø)`
src/awkward/operations/ak_argcartesian.py	`82.60% <ø> (ø)`
src/awkward/operations/ak_argcombinations.py	`92.30% <ø> (ø)`
src/awkward/operations/ak_argmax.py	`94.59% <ø> (ø)`
src/awkward/operations/ak_argmin.py	`94.59% <ø> (ø)`
src/awkward/operations/ak_argsort.py	`75.00% <ø> (ø)`
src/awkward/operations/ak_cartesian.py	`92.17% <ø> (ø)`
src/awkward/operations/ak_combinations.py	`88.46% <ø> (ø)`
src/awkward/operations/ak_concatenate.py	`85.86% <ø> (ø)`
... and 28 more

... and 3 files with indirect coverage changes

🚀 New features to boost your workflow:

❄️ Test Analytics: Detect flaky tests, report on failures, and find test suite problems.

github-actions · 2026-01-25T08:43:22Z

The documentation preview is ready to be viewed at http://preview.awkward-array.org.s3-website.us-east-1.amazonaws.com/PR3823

ikrommyd · 2026-01-25T08:57:57Z

cc @pfackeldey since you are much better at typing than me. I don’t know if this is the best though because AxisName already includes int and None right?

ianna

@X0708a - Thanks! It's a great start and IMHO the task of axis clarification is much bigger that adding it to one function.

Also, to make the reference manual actually useful for someone encountering AxisName, AxisTuple, or AxisMapping terms, they should linked to some explanation (btw, does it exist in docs?). The type aliases should be explained in terms of the functional role of them in the metadata.

I think, I would prefer consistency in the documentation and have one PR that clarifies axis in all high level functions that use it. It would be a much bigger job though :-)

src/awkward/operations/ak_moment.py

ikrommyd · 2026-01-26T08:56:09Z

FYI, named axes can only be strings. See

awkward/src/awkward/_namedaxis.py

Lines 137 to 163 in 680ab6e

    
           def _is_valid_named_axis(axis: AxisName) -> bool: 
        
               """ 
        
               Checks if the given axis is a valid named axis. A valid named axis is a hashable object that is not an integer or None. Currently it is restricted to strings. 
        
               Args: 
        
                   axis (AxisName): The axis to check. 
        
               Returns: 
        
                   bool: True if the axis is a valid named axis, False otherwise. 
        
               Examples: 
        
                   >>> _is_valid_named_axis("x") 
        
                   True 
        
                   >>> _is_valid_named_axis(1) 
        
                   False 
        
               """ 
        
               return ( 
        
                   # axis must be hashable 
        
                   isinstance(axis, AxisName) 
        
                   # ... but not an integer, otherwise we would confuse it with positional axis 
        
                   and not is_integer(axis) 
        
                   # we also prohibit None, which is reserved for wildcard 
        
                   and axis is not None 
        
                   # Let's only allow strings for now, in the future we can open up to more types 
        
                   # by removing the isinstance(axis, str) check. 
        
                   and isinstance(axis, str) 
        
               )

Therefore, I think for the tine being "int or None or string" is fine for the docstrings (and one can add a small oneliner that string is only in the case the axes are named).

Something like int or None or string and also add something like "If named axes are attached to the input(s), those can be used for the axis parameter as well (see also ...) and point them to the ak.with_named_axis function and the named axis user guide.
For all high-level functions of course.

aashirvad08 · 2026-01-26T19:35:46Z

Updated user-facing typing and docstrings across high-level reducers to
explicitly document support for named axes using strings, without exposing
internal typing aliases.

ikrommyd · 2026-01-26T20:04:46Z

There is no such commit that I see.

…ent-v2

aashirvad08 · 2026-01-27T05:04:55Z

Thanks for flagging this - the updates were on a different branch.
I’ve merged them into fix-axis-typing-ak-moment-v2 now; please let me know
if everything looks correct.

ikrommyd

You have introduced type hints in places where we don't have type hints and for only one parameter. You should only change docstrings here and not for just reducers but for all highlevel functions that have an axis parameter. Your docstring addition should probably point the users to the named axis user guide. No type hints should be introduced in places where there aren't any.

aashirvad08 · 2026-01-27T15:26:05Z

I’ve rebased onto the latest state of the branch and updated axis docstrings
across all high-level functions to document named-axis support and reference
the named-axis user guide.

aashirvad08 · 2026-01-29T05:34:20Z

I believe all requested changes have been addressed. Please let me know if there’s anything else I should adjust.

ikrommyd · 2026-01-29T09:42:06Z

The current changes are definitely not for all highlevel functions included under src/awkward/operations. I see only 7 at the moment. You should probably also not delete exising parts of the docstring. You are also not linking to the named axes docs. I would do something like this:

diff --git a/src/awkward/operations/ak_moment.py b/src/awkward/operations/ak_moment.py
index e8433a02..b9a16fe7 100644
--- a/src/awkward/operations/ak_moment.py
+++ b/src/awkward/operations/ak_moment.py
@@ -44,11 +44,15 @@ def moment(
             weight. Weighting values equally is the same as no weights;
             weighting some values higher increases the significance of those
             values. Weights can be zero or negative.
-        axis (None or int): If None, combine all values from the array into
+        axis (None or int or str): If None, combine all values from the array into
             a single scalar result; if an int, group by that axis: `0` is the
             outermost, `1` is the first level of nested lists, etc., and
             negative `axis` counts from the innermost: `-1` is the innermost,
-            `-2` is the next level up, etc.
+            `-2` is the next level up, etc. If a str, it is interpreted as 
+            the name of the axis which maps to an int if named axes are present.
+            Named axes are attached to an array using
+            #ak.with_named_axis and removed with #ak.without_named_axis; also
+            see the [Named axes user guide](../../user-guide/how-to-array-properties-named-axis.html).
         keepdims (bool): If False, this function decreases the number of
             dimensions by 1; if True, the output values are wrapped in a new
             length-1 dimension so that the result of this operation may be

but it needs to be done for ALL highlevel functions similarly.

aashirvad08 · 2026-01-29T10:26:10Z

Thanks for the concrete example, that’s very helpful.

Understood — I’ll update the axis docstrings following this pattern (preserving
existing text, extending it to cover named axes, and linking to the named-axis
user guide), and apply it consistently across all high-level functions under
src/awkward/operations that accept an axis parameter.

aashirvad08 · 2026-01-31T19:57:54Z

I’ve now updated all high-level operations to consistently document named-axis support for axis, without removing existing docstring content, and added a link to the named-axis user guide.
Do let me know if anything else is required.

ianna

@X0708a - thank for addressing the comments! The axis doc seems out of place. Please remove it. Also- have you checked all other high level functions that take axis parameter? Btw, the pr title doesn’t correspond to its content anymore 😀

docs/user-guide/axis.md

src/awkward/operations/ak_argmin.py

src/awkward/operations/ak_var.py

src/awkward/operations/ak_prod.py

aashirvad08 · 2026-02-03T19:09:49Z

Thanks for the careful review and the ripgrep lists was very helpful.
I’ve now addressed all the remaining points raised in the review.
Please let me know if anything else needs adjustment.

src/awkward/operations/ak_sum.py

ikrommyd · 2026-02-03T22:21:04Z

I only looked at the changed files now and I think you now have touched all highlevel functions. Please someone else double check me too. This is the complete ripgrep

❯ rg "axis \(" src/awkward/operations | sort
src/awkward/operations/ak_all.py:        axis (None or int): If None, combine all values from the array into
src/awkward/operations/ak_any.py:        axis (None or int): If None, combine all values from the array into
src/awkward/operations/ak_argcartesian.py:        axis (int): The dimension at which this operation is applied. The
src/awkward/operations/ak_argcombinations.py:        axis (int): The dimension at which this operation is applied. The
src/awkward/operations/ak_argmax.py:        axis (None or int): If None, combine all values from the array into
src/awkward/operations/ak_argmax.py:        axis (None or int): If None, combine all values from the array into
src/awkward/operations/ak_argmin.py:        axis (None or int): If None, combine all values from the array into
src/awkward/operations/ak_argmin.py:        axis (None or int): If None, combine all values from the array into
src/awkward/operations/ak_argsort.py:        axis (int): The dimension at which this operation is applied. The
src/awkward/operations/ak_broadcast_arrays.py:    explicitly adding a new axis (reshape to add a dimension of length 1)
src/awkward/operations/ak_cartesian.py:        axis (int): The dimension at which this operation is applied. The
src/awkward/operations/ak_combinations.py:        axis (int): The dimension at which this operation is applied. The
src/awkward/operations/ak_concatenate.py:        axis (int): The dimension at which this operation is applied. The
src/awkward/operations/ak_corr.py:        axis (None or int): If None, combine all values from the array into
src/awkward/operations/ak_count_nonzero.py:        axis (None or int): If None, combine all values from the array into
src/awkward/operations/ak_count.py:        axis (None or int): If None, combine all values from the array into
src/awkward/operations/ak_covar.py:        axis (None or int): If None, combine all values from the array into
src/awkward/operations/ak_drop_none.py:        axis (None or int): If None, the operation drops Nones at all levels of
src/awkward/operations/ak_fill_none.py:        axis (None or int): If None, replace all None values in the array
src/awkward/operations/ak_firsts.py:        axis (int): The dimension at which this operation is applied. The
src/awkward/operations/ak_flatten.py:        axis (None or int): If None, the operation flattens all levels of
src/awkward/operations/ak_from_regular.py:        axis (int or None): The dimension at which this operation is applied.
src/awkward/operations/ak_is_none.py:        axis (int): The dimension at which this operation is applied. The
src/awkward/operations/ak_linear_fit.py:        axis (None or int): If None, combine all values from the array into
src/awkward/operations/ak_local_index.py:        axis (int): The dimension at which this operation is applied. The
src/awkward/operations/ak_max.py:        axis (None or int): If None, combine all values from the array into
src/awkward/operations/ak_max.py:        axis (None or int): If None, combine all values from the array into
src/awkward/operations/ak_mean.py:        axis (None or int): If None, combine all values from the array into
src/awkward/operations/ak_mean.py:        axis (None or int): If None, combine all values from the array into
src/awkward/operations/ak_merge_option_of_records.py:        axis (int): The dimension at which this operation is applied.
src/awkward/operations/ak_merge_union_of_records.py:        axis (int): The dimension at which this operation is applied.
src/awkward/operations/ak_min.py:        axis (None or int): If None, combine all values from the array into
src/awkward/operations/ak_min.py:        axis (None or int): If None, combine all values from the array into
src/awkward/operations/ak_moment.py:        axis (None or int): If None, combine all values from the array into
src/awkward/operations/ak_num.py:        axis (int): The dimension at which this operation is applied. The
src/awkward/operations/ak_pad_none.py:        axis (int): The dimension at which this operation is applied. The
src/awkward/operations/ak_prod.py:        axis (None or int): If None, combine all values from the array into
src/awkward/operations/ak_prod.py:        axis (None or int): If None, combine all values from the array into
src/awkward/operations/ak_ptp.py:        axis (None or int): If None, combine all values from the array into
src/awkward/operations/ak_singletons.py:        axis (int): The dimension at which this operation is applied. The
src/awkward/operations/ak_softmax.py:        axis (None or int): If None, combine all values from the array into
src/awkward/operations/ak_sort.py:        axis (int): The dimension at which this operation is applied. The
src/awkward/operations/ak_std.py:        axis (None or int): If None, combine all values from the array into
src/awkward/operations/ak_std.py:        axis (None or int): If None, combine all values from the array into
src/awkward/operations/ak_sum.py:        axis (None or int): If None, combine all values from the array into
src/awkward/operations/ak_sum.py:        axis (None or int): If None, combine all values from the array into
src/awkward/operations/ak_to_regular.py:        axis (int or None): The dimension at which this operation is applied.
src/awkward/operations/ak_unflatten.py:        axis (int): The dimension at which this operation is applied. The
src/awkward/operations/ak_var.py:        axis (None or int): If None, combine all values from the array into
src/awkward/operations/ak_var.py:        axis (None or int): If None, combine all values from the array into

ikrommyd · 2026-02-03T22:22:18Z

I noticed that you correctly did not change the ak.to/from_regular docstrings as those functions don't support named axes. But I don't know why that is and it seems to me that they probably should support named axes. Any reason @pfackeldey you might remember?

src/awkward/operations/ak_concatenate.py

ikrommyd

I've pushed a couple of fixes. It looks good to me now but I may still have missed some inconsistency. @ariostas could you re-check too?

ikrommyd

I think it should be good now? I don't see anything else in the docstrings but please someone else have a look too.

Copilot

Pull request overview

This PR updates documentation for the axis parameter across many high-level operations to document support for named axes (string axis names). The main change removes the problematic AxisName type annotation from ak.moment (which incorrectly allowed None as a hashable) and updates docstrings across 33 operation files to clarify that axis can be None, int, or str.

Changes:

Removed axis: AxisName = None type annotation from ak.moment, replacing it with axis=None to avoid type hint issues
Updated docstrings across 33 operations to document that axis accepts None, int, or str values
Added consistent documentation blocks explaining named axis support with links to the user guide

Reviewed changes

Copilot reviewed 38 out of 38 changed files in this pull request and generated 3 comments.

Show a summary per file

File	Description
src/awkward/operations/ak_moment.py	Removed problematic `AxisName` type annotation from axis parameter; updated docstring
src/awkward/operations/ak_var.py	Updated axis docstring to document named-axis support for `var` and `nanvar`
src/awkward/operations/ak_unflatten.py	Updated axis docstring to document named-axis support
src/awkward/operations/ak_sum.py	Updated axis docstring to document named-axis support for `sum` and `nansum`
src/awkward/operations/ak_std.py	Updated axis docstring to document named-axis support for `std` and `nanstd`
src/awkward/operations/ak_sort.py	Updated axis docstring to document named-axis support
src/awkward/operations/ak_softmax.py	Updated and restructured axis docstring to document named-axis support
src/awkward/operations/ak_singletons.py	Updated axis docstring to document named-axis support
src/awkward/operations/ak_ptp.py	Updated axis docstring to document named-axis support
src/awkward/operations/ak_prod.py	Updated axis docstring to document named-axis support for `prod` and `nanprod`
src/awkward/operations/ak_pad_none.py	Updated axis docstring to document named-axis support
src/awkward/operations/ak_num.py	Updated axis docstring to document named-axis support
src/awkward/operations/ak_min.py	Updated axis docstring to document named-axis support for `min` and `nanmin`
src/awkward/operations/ak_merge_union_of_records.py	Updated axis docstring to document named-axis support
src/awkward/operations/ak_merge_option_of_records.py	Updated axis docstring to document named-axis support
src/awkward/operations/ak_mean.py	Updated axis docstring to document named-axis support for `mean` and `nanmean`
src/awkward/operations/ak_max.py	Updated axis docstring to document named-axis support for `max` and `nanmax`
src/awkward/operations/ak_local_index.py	Updated axis docstring to document named-axis support
src/awkward/operations/ak_linear_fit.py	Updated axis docstring to document named-axis support
src/awkward/operations/ak_is_none.py	Updated axis docstring to document named-axis support
src/awkward/operations/ak_flatten.py	Updated axis docstring to document named-axis support
src/awkward/operations/ak_firsts.py	Updated axis docstring to document named-axis support
src/awkward/operations/ak_fill_none.py	Updated axis docstring to document named-axis support
src/awkward/operations/ak_drop_none.py	Updated axis docstring to document named-axis support
src/awkward/operations/ak_covar.py	Updated axis docstring to document named-axis support
src/awkward/operations/ak_count_nonzero.py	Updated axis docstring to document named-axis support
src/awkward/operations/ak_count.py	Updated axis docstring to document named-axis support
src/awkward/operations/ak_corr.py	Updated axis docstring to document named-axis support
src/awkward/operations/ak_concatenate.py	Updated axis docstring to document named-axis support
src/awkward/operations/ak_combinations.py	Updated axis docstring to document named-axis support
src/awkward/operations/ak_cartesian.py	Updated axis docstring to document named-axis support
src/awkward/operations/ak_argsort.py	Updated axis docstring to document named-axis support
src/awkward/operations/ak_argmin.py	Updated axis docstring to document named-axis support for `argmin` and `nanargmin`
src/awkward/operations/ak_argmax.py	Updated axis docstring to document named-axis support for `argmax` and `nanargmax`
src/awkward/operations/ak_argcombinations.py	Updated axis docstring to document named-axis support
src/awkward/operations/ak_argcartesian.py	Updated axis docstring to document named-axis support
src/awkward/operations/ak_any.py	Updated axis docstring to document named-axis support
src/awkward/operations/ak_all.py	Updated axis docstring to document named-axis support

💡 Add Copilot custom instructions for smarter, more guided reviews. Learn how to get started.

src/awkward/operations/ak_moment.py

src/awkward/operations/ak_fill_none.py

src/awkward/operations/ak_argmin.py

Co-authored-by: Copilot <175728472+Copilot@users.noreply.github.com>

ariostas

Seems good to me now

Clarify axis typing and documentation for ak.moment

aa90642

aashirvad08 mentioned this pull request Jan 25, 2026

Clarify axis typing and documentation for ak.moment #3822

Closed

ianna requested a review from pfackeldey January 25, 2026 13:23

ianna requested changes Jan 25, 2026

View reviewed changes

src/awkward/operations/ak_moment.py Outdated Show resolved Hide resolved

ikrommyd mentioned this pull request Jan 26, 2026

Docs: explain axis semantics and axis-related types #3825

Closed

aashirvad08 added 4 commits January 26, 2026 23:53

docs: clarify axis parameter supports named axes using strings

e71b3d3

docs: clarify axis parameter supports named axes in ak.sum

e65591c

docs: clarify axis parameter supports named axes in ak.sum

73fc037

docs: clarify axis parameter supports named axes across reducers

2534cbe

aashirvad08 and others added 2 commits January 27, 2026 10:32

Merge branch 'axis-docstrings-named-axes' into fix-axis-typing-ak-mom…

f2ba860

…ent-v2

style: pre-commit fixes

4818eec

ikrommyd requested changes Jan 27, 2026

View reviewed changes

docs: clarify axis docstrings and reference named-axis guide

2a0cb5b

aashirvad08 and others added 2 commits February 1, 2026 01:24

docs: document named-axis support for axis in high-level operations

693af3b

style: pre-commit fixes

17436ff

ianna requested changes Feb 1, 2026

View reviewed changes

docs/user-guide/axis.md Outdated Show resolved Hide resolved

docs: remove draft axis user guide from this PR

2823898

aashirvad08 changed the title ~~docs: clarify axis typing and documentation for ak.moment~~ docs: document named-axis support for axis in high-level operations Feb 1, 2026

ikrommyd reviewed Feb 3, 2026

View reviewed changes

src/awkward/operations/ak_argmin.py Outdated Show resolved Hide resolved

Merge branch 'main' into fix-axis-typing-ak-moment-v2

28941f4

ariostas reviewed Feb 3, 2026

View reviewed changes

src/awkward/operations/ak_var.py Outdated Show resolved Hide resolved

src/awkward/operations/ak_prod.py Outdated Show resolved Hide resolved

aashirvad08 and others added 2 commits February 4, 2026 00:36

docs: fix remaining axis docstrings and formatting issues

f27cdf1

style: pre-commit fixes

384ed0b

ariostas reviewed Feb 3, 2026

View reviewed changes

src/awkward/operations/ak_sum.py Show resolved Hide resolved

docs: restore array argument in ak_sum docstring

dbd92da

ikrommyd reviewed Feb 3, 2026

View reviewed changes

src/awkward/operations/ak_concatenate.py Outdated Show resolved Hide resolved

ikrommyd added 9 commits February 3, 2026 16:30

fix locations where named axes guide is not linked

655be27

undo code change

03538d9

fix nanvar

bc4d657

various fixes

c98ec24

Merge branch 'main' into fix-axis-typing-ak-moment-v2

2cafce9

more fixes

9ce0eb7

missing bracket

f4d40e5

fix firsts

f4f2d42

fix num

4f80752

ikrommyd reviewed Feb 3, 2026

View reviewed changes

ikrommyd approved these changes Feb 5, 2026

View reviewed changes

ariostas requested a review from Copilot February 5, 2026 18:17

Copilot started reviewing on behalf of ariostas February 5, 2026 18:17 View session

Copilot AI reviewed Feb 5, 2026

View reviewed changes

src/awkward/operations/ak_moment.py Show resolved Hide resolved

src/awkward/operations/ak_fill_none.py Outdated Show resolved Hide resolved

src/awkward/operations/ak_argmin.py Outdated Show resolved Hide resolved

ariostas and others added 2 commits February 5, 2026 13:24

Update src/awkward/operations/ak_fill_none.py

bea3e89

Co-authored-by: Copilot <175728472+Copilot@users.noreply.github.com>

Update src/awkward/operations/ak_argmin.py

a83d900

Co-authored-by: Copilot <175728472+Copilot@users.noreply.github.com>

ariostas approved these changes Feb 5, 2026

View reviewed changes

ikrommyd merged commit 1e1ca0c into scikit-hep:main Feb 5, 2026
37 of 39 checks passed

ikrommyd mentioned this pull request Feb 5, 2026

Inaccurate typing / documentation of axis parameters #3488

Closed

Conversation

aashirvad08 commented Jan 25, 2026 • edited by ikrommyd Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Uh oh!

codecov bot commented Jan 25, 2026 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Codecov Report

Uh oh!

github-actions bot commented Jan 25, 2026

Uh oh!

ikrommyd commented Jan 25, 2026 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Uh oh!

ianna left a comment

Choose a reason for hiding this comment

Uh oh!

Uh oh!

ikrommyd commented Jan 26, 2026 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Uh oh!

aashirvad08 commented Jan 26, 2026

Uh oh!

ikrommyd commented Jan 26, 2026 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Uh oh!

aashirvad08 commented Jan 27, 2026

Uh oh!

ikrommyd left a comment • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Choose a reason for hiding this comment

Uh oh!

aashirvad08 commented Jan 27, 2026

Uh oh!

aashirvad08 commented Jan 29, 2026

Uh oh!

ikrommyd commented Jan 29, 2026

Uh oh!

aashirvad08 commented Jan 29, 2026

Uh oh!

aashirvad08 commented Jan 31, 2026

Uh oh!

ianna left a comment

Choose a reason for hiding this comment

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

aashirvad08 commented Feb 3, 2026

Uh oh!

Uh oh!

ikrommyd commented Feb 3, 2026

Uh oh!

ikrommyd commented Feb 3, 2026

Uh oh!

Uh oh!

ikrommyd left a comment

Choose a reason for hiding this comment

Uh oh!

ikrommyd left a comment

Choose a reason for hiding this comment

Uh oh!

Copilot AI left a comment

Choose a reason for hiding this comment

Pull request overview

Reviewed changes

Uh oh!

Uh oh!

Uh oh!

Uh oh!

ariostas left a comment

Choose a reason for hiding this comment

Uh oh!

Uh oh!

Reviewers

Assignees

Labels

Projects

Milestone

Development

aashirvad08 commented Jan 25, 2026 •

edited by ikrommyd

Loading

codecov bot commented Jan 25, 2026 •

edited

Loading

ikrommyd commented Jan 25, 2026 •

edited

Loading

ikrommyd commented Jan 26, 2026 •

edited

Loading

ikrommyd commented Jan 26, 2026 •

edited

Loading

ikrommyd left a comment •

edited

Loading