MAINT: Speed up normalize_axis_tuple by about 30% #11518
Conversation
This is used in `np.moveaxis`, which is relatively slow. Timings change as follows: ``` import numpy as np from numpy.core.numeric import normalize_axis_tuple %timeit normalize_axis_tuple([0, 1, 2], 4, 'parrot') a = np.zeros((10,20,30)) %timeit np.moveaxis(a, [0,1,2], [2,1,0]) 100000 loops, best of 3: 5.56 µs per loop -> 4.17 µs ```
numpy/core/numeric.py
Outdated
| axis = tuple(axis) | ||
| axis = tuple(normalize_axis_index(ax, ndim, argname) for ax in axis) | ||
| # Speed-up most common cases. | ||
| if not isinstance(axis, (list, tuple)): |
There was a problem hiding this comment.
Would be more correct to check type(axis) in (list, tuple), since subclasses might override __index__.
I was tempted to change it to just type(axis) is tuple, to encourage people to use tuple rather than list axis arguments....
There was a problem hiding this comment.
I think if someone subclasses list or tuple, one should reasonably expect that the object can be iterated over, not also treated as an integer! (I'd actually prefer to just test for something being iterable, but things like np.iterable are way slower than the isinstance check.)
As for tuple - I can at least change the order!
There was a problem hiding this comment.
Reasonable or not, I'd prefer that we don't introduce an unnecessary behavior change now.
type(axis) in (tuple, list) sounds like a good approach to me
| except TypeError: | ||
| pass | ||
| # Going via an iterator directly is slower than via list comprehension. | ||
| axis = tuple([normalize_axis_index(ax, ndim, argname) for ax in axis]) |
There was a problem hiding this comment.
I don't understand why this is the case...
There was a problem hiding this comment.
I think list comprehension is special-cased, not constructing an iterator but an actual for-loop with known size.
|
@eric-wieser - since you wrote the function, maybe it is easiest for you to just make any changes you want? I don't really like what I have here but don't have the time to think of something better. |
|
Updated in place. I think I only extracted the function from its call site, rather than inventing it |
|
OK, let's get this in - it is a speed improvement even if not super elegant. |
|
Thanks for identifying a slowdown here - I'll be curious to see the benchmarks when they update |
This is used in
np.moveaxis, which is relatively slow. Timings change as follows:Should add that I'm not really all that happy with the present solution, but I cannot think of a better one that wouldn't either impact the iterable or the scalar case.