MAINT: Avoid moveaxis overhead in median. #18324
Conversation
|
What does this do for high dimensional arrays, say I can't say I like this, but I can live with it. If we want to do this, do we have a benchmark? |
|
fwiw Do you want something like the following? diff --git i/benchmarks/benchmarks/bench_lib.py w/benchmarks/benchmarks/bench_lib.py
index c22ceaa5e..b86b2f6e9 100644
--- i/benchmarks/benchmarks/bench_lib.py
+++ w/benchmarks/benchmarks/bench_lib.py
@@ -6,6 +6,23 @@ from .common import Benchmark
import numpy as np
+class Median(Benchmark):
+ """Benchmarks for `numpy.median`"""
+
+ param_names = ["array_size", "axis"]
+ params = [
+ [(1001,), (10000, 20)],
+ [None, -1],
+ ]
+
+ def setup(self, array_size, axis):
+ np.random.seed(123)
+ self.array = np.random.random(array_size)
+
+ def time_median(self, array_size, axis):
+ np.median(self.array, axis=axis)
+
+
class Pad(Benchmark):
"""Benchmarks for `numpy.pad`. |
|
Yeah, was thinking of something like that. Seems we already have benchmarks including "small arrays" (500 elements) in Anyway, if no-one else has an opinion... It feels hackish to me, but I guess we should just roll with it and feel free to change it later (hopefully we would notice if there is a serious performance regression with the benchmarks). I am wondering now, if I got the axes wrong, and it should have been |
This change speeds up taking the median of 1001 floats by ~20%, as measured by `python -mtimeit -s 'import numpy as np; x = np.random.randn(1001)' -- 'np.median(x)'`
I agree making moveaxis/asarray faster would be better, but that's beyond my pay grade :)
Added the benchmarks there.
timeit doesn't show a difference in this case either. |
|
Are you suggesting that |
|
@Qiyu8 take is pretty much never good performance wise, this is absolutely an exception, and even here its more that moveaxis should be faster probably. |
|
I did try to change the polynomial calls to moveaxis but t
8000
here there was no performance improvement. |
|
A not-so-small part of it may also be that Moving Speeding up |
|
Slightly faster still (but less elegant to me) is Ideally, |
|
I can go with @mhvk's solution if the rest agree (it seems best, except for the slight ugliness of the code needed to set up the right index tuple). Alternatively (but more complex), given the existence of the same problem in #18203, perhaps a solution would be to fold nan-checking into PyArray_Partition and to have a variant of it ( |
|
I am ok with it, its ugly, but maybe we should add a brief comment that it is a micro-opt for moveaxes. The doc of the function suggests multiple axes are allowed. that seems simply not true, but lets fix that? I guess we do have a test for things like: I am not sure how easy it is to check for NaN during partition, I think if anything, it might be more useful to create an "optimization" gufunc "all_finite()" or so. Applying that to the upper half is in theory a bit faster probably. EDIT: To be a bit more assertive, the list trick is fairly common also in other functions I think. Maybe |
|
Definitely need to check the multiple-axis case... To me, |
|
We don't actually need to support multiple axis here because median goes through _ureduce for that. In fact the docstring of _median_nancheck is wrong, it doesnt support axis being a sequence of int (the call to moveaxis() would fail). See #18379. |
|
OK, thanks @anntzer. I agree Going to put this in, although I am would be more than happy to pull it out again for a better alternative ;). |
This change speeds up taking the median of 1001 floats by ~20%, as
measured by
python -mtimeit -s 'import numpy as np; x = np.random.randn(1001)' -- 'np.median(x)'(See also #18298, #18203.)