Repository: numpy/numpy
Pull request: DOC: Improve description of ``axis`` parameter for ``np.median``
This PR includes an updated docstring for `numpy.median` which clarifies the functionality of the extended axis feature, introduced in 1.9.0. 

Specifically, the description of the docstring is changed and an example is added. These help to clarify if a sequence of axes is passed, that the array is flattened along the given axes and _then_ the median is computed over the resulting flattened axis. Without this change, the functionality of passing multiple axes could be misunderstood as computing medians sequentially over the given axes (without flattening), see #25174.

Resolves #25174

Problem statement:
DOC: Unexpected values from `np.median` when passing a sequence of ints for `axis=` 
### Describe the issue:

Calling `np.median` and passing a sequence of ints to the `axis=` parameter does not produce the same (expected) results as sequentially calling median on each axis of the sequence. My expectation is that these two methods of computing a median would produce the same result.

### Reproduce the code example:

```python
import numpy as np
a = [[[1,2],[3,4]],[[5,6],[8,9]]]
np.median(np.median(np.median(a, axis=0), axis=0), axis=0)
#> 4.75 
np.median(a, axis=[0,1,2])
#> 4.5
```


### Error message:

_No response_

### Runtime information:

```
>>> import sys, numpy; print(numpy.__version__); print(sys.version)
1.26.1
3.11.2 (tags/v3.11.2:878ead1, Feb  7 2023, 16:38:35) [MSC v.1934 64 bit (AMD64)]
>>> print(numpy.show_runtime())
WARNING: `threadpoolctl` not found in system! Install it by `pip install threadpoolctl`. Once installed, try `np.show_runtime` again for more detailed build information
[{'numpy_version': '1.26.1',
  'python': '3.11.2 (tags/v3.11.2:878ead1, Feb  7 2023, 16:38:35) [MSC v.1934 '
            '64 bit (AMD64)]',
  'uname': uname_result(system='Windows', node='FRST-FRM-2231E', release='10', version='10.0.19045', machine='AMD64')},
 {'simd_extensions': {'baseline': ['SSE', 'SSE2', 'SSE3'],
                      'found': ['SSSE3', 'SSE41', 'POPCNT', 'SSE42', 'AVX'],
                      'not_found': ['F16C',
                                    'FMA3',
                                    'AVX2',
                                    'AVX512F',
                                    'AVX512CD',
                                    'AVX512_SKX',
                                    'AVX512_CLX',
                                    'AVX512_CNL',
                                    'AVX512_ICL']}}]
None
```

### Context for the issue:

_No response_

Hints from issue discussion:
`axis=[0, 1, 2]` means that those axes are flattened, and the median is taken over the resulting flattened axis. This is not the same as taking the median of medians of medians along each axis sequentially.

The code is behaving as intended, but the docstring doesn't actually describe the semantics in thorough-enough detail.
Thank you for the clarification. I could submit a PR for this, if that works?
[Yes, please!](https://numpy.org/doc/1.26/dev/howto-docs.html#contributing-fixes)

For reference, the general `axis=seq_of_ints` semantics were introduced to numpy in [1.7.0](https://numpy.org/doc/1.26/release/1.7.0-notes.html#reduction-ufuncs-generalize-axis-parameter); you can see how it's documented in [`np.sum`](https://numpy.org/doc/1.26/reference/generated/numpy.sum.html) for instance. It was added to `np.median` in [1.9.0](https://numpy.org/doc/1.26/release/1.9.0-notes.html#generalized-axis-support-for-median-and-percentile), but the release notes and PR (#3908) are not especially enlightening.

Commit message: DOC: add sequence ex to docstring of numpy.median

The docstring for numpy.median does not contain an example for passing a
sequence of ints to the axis parameter. Including an example could help
clarify the functionality of the extended axis feature. This channge
includes a new example demonstrating the median with extended axis.

Patch to review:
diff --git a/numpy/lib/_function_base_impl.py b/numpy/lib/_function_base_impl.py
index 9693293bad7c..049db6698f78 100644
--- a/numpy/lib/_function_base_impl.py
+++ b/numpy/lib/_function_base_impl.py
@@ -3778,9 +3778,15 @@ def median(a, axis=None, out=None, overwrite_input=False, keepdims=False):
     a : array_like
         Input array or object that can be converted to an array.
     axis : {int, sequence of int, None}, optional
-        Axis or axes along which the medians are computed. The default
-        is to compute the median along a flattened version of the array.
-        A sequence of axes is supported since version 1.9.0.
+        Axis or axes along which the medians are computed. The default,
+        axis=None, will compute the median along a flattened version of
+        the array.
+        
+        .. versionadded:: 1.9.0
+
+        If a sequence of axes, the array is first flattened along the
+        given axes then the median is computed along the resulting
+        flattened axis.
     out : ndarray, optional
         Alternative output array in which to place the result. It must
         have the same shape and buffer length as the expected output,
@@ -3832,6 +3838,8 @@ def median(a, axis=None, out=None, overwrite_input=False, keepdims=False):
     array([6.5, 4.5, 2.5])
     >>> np.median(a, axis=1)
     array([7.,  2.])
+    >>> np.median(b, axis=(0,1), overwrite_input=True)
+    3.5
     >>> m = np.median(a, axis=0)
     >>> out = np.zeros_like(m)
     >>> np.median(a, axis=0, out=m)


Code review comments:
diff_hunk: @@ -3832,6 +3838,8 @@ def median(a, axis=None, out=None, overwrite_input=False, keepdims=False):
     array([6.5, 4.5, 2.5])
     >>> np.median(a, axis=1)
     array([7.,  2.])
+    >>> np.median(b, axis=(0,1), overwrite_input=True)
+    3.5
path: numpy/lib/_function_base_impl.py
text: @user2:
These examples in the docstrings should be runnable. Using `b` here when it hasn't been defined above doesn't make sense. I would also not use `overwrite_input`, since that confuses things. 

In principle doctests should catch issues like `b` not being defined in the docstring yet, but I'm not sure we run doctests on the full numpy codebase...

Instead, I think it makes more sense to just rewrite all of these examples to use a small 3D array, so this example where you take the median successively over two dimensions makes sense. Something like:

```python
>>> import numpy as np
>>> a = np.arange(27)
>>> a.shape = (3, 3, 3)
>>> a
array([[[ 0,  1,  2],
        [ 3,  4,  5],
        [ 6,  7,  8]],

       [[ 9, 10, 11],
        [12, 13, 14],
        [15, 16, 17]],

       [[18, 19, 20],
        [21, 22, 23],
        [24, 25, 26]]])
>>> np.median(a)
np.float64(13.0)
>>> np.median(a, axis=0)
array([[ 9., 10., 11.],
       [12., 13., 14.],
       [15., 16., 17.]])
>>> np.median(a, axis=(0, 1))
array([12., 13., 14.])
```

And then if you didn't feel like editing the rest of this, you could redefine a to its original value below this example.

Just a thought, though! There are probably other ways to edit this that would make sense.

Also it would be nice if you could make sure this docstring is actually fully runnable and updated, since it looks like it hasn't been updated for the fallout from NEP 50 (ping @user1 - there's probably a lot of other docstrings that need to be updated right?).

@user2:
Err, not NEP 50, I guess just the new reprs for scalars?

@user1:
Yeah, that is on my not-done list.  I had always hoped to use https://github.com/scientific-python/pytest-doctestplus/pull/227 although that is a bit tedious until we better support doctestplus.

OTOH, a few failures (and misses) shouldn't matter too much.

At this time, I think just don't worry about whether or not you write `np.float64(1.0)` or `1.0`, keeping things consistent may be better...