numpy ndindex is slower than itertools alternative #28921
Description
Proposed new feature or change:
I don't know if I am missing something, but numpy ndindex is underperforming over all input sizes here (I tried different ndims as well):
from itertools import product
from numpy import ndindex as np_ndindex
def consume(x):
for _ in x:
pass
def py_ndindex(*shape):
"""Yield tuples of indices for an array of given shape."""
return product(*(range(s) for s in shape))
%timeit -n 1 -r 1 consume(py_ndindex(50, 60, 90)) # 9.97 ms ± 0 ns per loop (mean ± std. dev. of 1 run, 1 loop each)
%timeit -n 1 -r 1 consume(np_ndindex(50, 60, 90)) # 98.4 ms ± 0 ns per loop (mean ± std. dev. of 1 run, 1 loop each)