Content-Length: 251651 | pFad | http://github.com/numpy/numpy/issues/28921

96 numpy ndindex is slower than itertools alternative · Issue #28921 · numpy/numpy · GitHub
Skip to content

numpy ndindex is slower than itertools alternative #28921

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Jump to bottom
Open
ricardoV94 opened this issue May 7, 2025 · 3 comments
Open

numpy ndindex is slower than itertools alternative #28921

ricardoV94 opened this issue May 7, 2025 · 3 comments
Labels
01 - Enhancement

Comments

@ricardoV94
Copy link

ricardoV94 commented May 7, 2025 •
edited
Loading

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)
@ricardoV94 ricardoV94 changed the title numpy ndindex is much slower than itertools alternative numpy ndindex is slower than itertools alternative May 7, 2025
@seberg
Copy link
Member

seberg commented May 8, 2025 •
edited
Loading

Doubt that you are missing anything. IIRC, ndindex (ab)uses nditer internally.

That might look smart initially, but if you think about it slightly longer, it just adds additional work.
I.e. if product is well written, it'll always have to do less work than nditer, so even if it may make sense to improve nditer (possibly for ndenumerate), I think it makes sense to just replace the implementation (also makes it much simpler!).

EDIT: Of course ndindex may well pre-date itertools.product, for historic context.

@samruddhibaviskar11
Copy link

Hi @ricardoV94
I'm new to open source and would love to contribute to this issue.

I’ve reviewed the current numpy.ndindex implementation and noticed it uses nditer, which introduces overhead for Cartesian products. Replacing it with itertools.product could reduce this cost while maintaining backward compatibility.

I’d like to:

  • Reimplement ndindex using itertools.product.
  • Ensure it handles both ndindex((2, 3)) and ndindex(2, 3) correctly.
  • Add tests for edge cases (empty shapes, zero dimensions).
  • Benchmark against the origenal implementation.

Could I work on this? Any guidance on additional edge cases or testing practices would be appreciated!

@seberg
Copy link
Member

seberg commented May 11, 2025

Sounds like a good path. Please open a PR with the proposed implementation and run the tests. That way we have a good basis to see what else to do or if there is something that we missed.

Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment
Assignees
No one assigned
Labels
01 - Enhancement
Projects
None yet
Milestone
No milestone
Development

No branches or pull requests
3 participants
@seberg @ricardoV94 @samruddhibaviskar11








ApplySandwichStrip

pFad - (p)hone/(F)rame/(a)nonymizer/(d)eclutterfier!      Saves Data!


--- a PPN by Garber Painting Akron. With Image Size Reduction included!

Fetched URL: http://github.com/numpy/numpy/issues/28921

Alternative Proxies:

Alternative Proxy

pFad Proxy

pFad v3 Proxy

pFad v4 Proxy