ENH: Use itertools.product for ndindex to improve performance #29165

samruddhibaviskar11 · 2025-06-10T21:51:09Z

Summary

This PR rewrites numpy.ndindex to use itertools.product instead of nditer, significantly improving both runtime performance and memory efficiency while preserving full backward compatibility.

It addresses Issue #28921: ndindex is slower than itertools alternative.

Motivation

The current implementation of ndindex uses nditer, which is flexible but adds unnecessary overhead when generating Cartesian product indices. Since itertools.product is a well-optimized C implementation for this exact pattern, switching to it results in a notable performance improvement.

Changes Made

Rewrote the ndindex.__init__ and __next__ methods in numpy/lib/_index_tricks_impl.py to utilize itertools.product.
Added a comprehensive set of new unit tests in numpy/lib/tests/test_index_tricks.py to cover various scenarios (e.g., empty/zero dimensions, non-integer dimensions, iterator independence, stop iteration behavior, large dimensions, and consistency with ndenumerate).
Introduced benchmarks/benchmarks/bench_ndindex.py to provide a clear performance comparison.
Updated the docstring for ndindex in _index_tricks_impl.py to include a Notes section explaining the implementation change and its benefits.

Benchmark Results

Benchmarks were added in benchmarks/benchmarks/bench_ndindex.py. Here is a summary:

==================================================================
                NumPy ndindex Performance Benchmark               
==================================================================

Shape            | Elements | Legacy (ms) | Optimized (ms) | Speedup
-----------------|----------|-------------|----------------|---------
(10, 10)         | 100      | 0.89        | 0.02           | 36.9×
(20, 20)         | 400      | 0.31        | 0.06           | 5.6×
(50, 50)         | 2,500    | 1.59        | 0.23           | 6.9×
(10, 10, 10)     | 1,000    | 0.84        | 0.14           | 5.9×
(20, 30, 40)     | 24,000   | 12.71       | 2.10           | 6.1×
(50, 60, 90)     | 270,000  | 130.78      | 25.68          | 5.1×

# => Average speedup: ~11.1×

Memory Usage Test:

Testing shape (100, 100) (10000 total elements)
Iterator creation time: 0.153 ms
Time to consume first 1000 elements: 0.615 ms
--> Memory efficient: Only generates indices as needed

SUCCESS! Average performance improvement: 9.8x
This demonstrates excellent performance improvement.

Test Cases & Functional Correctness

The following new tests were added to ensure correctness across a variety of use cases:

Empty shapes and zero dimensions
StopIteration behavior
Non-integer dimensions raise TypeError
Argument consistency (ndindex((2, 3)) vs ndindex(2, 3))
Compatibility with ndenumerate
Independent iterator instances
Large dimensions (e.g., (1000, 1000))
Empty iterators for zero-length dimensions
Performance regression sanity checks (<1s for large shapes)
Multidimensional correctness

Tests live in numpy/lib/tests/test_index_tricks.py.

Environment/Testing Notes:

✅ The optimized implementation was developed and tested locally in a standalone file (ndindex_optimized.py), separate from the NumPy codebase.

✅ A full set of unit tests were written and run successfully in isolation using pytest, validating correctness across all intended behaviors.

🔁 When attempting to run the full NumPy test suite inside a GitHub Codespace (editable install with pip install -e . and meson/ninja), I encountered the following persistent import error:

ImportError: cannot import name 'version' from partially initialized module 'numpy'

🔍 This appears to be a build tooling issue, not a problem with the PR's logic or correctness.

📦 Submitting this PR in hopes that GitHub Actions CI will successfully build and test the full suite, or maintainers can help identify the issue with the local dev environment

Docstring Update

A Notes section has been added to the ndindex docstring to reflect the new implementation:

Notes
-----
As of NumPy X.Y.Z, `ndindex` is implemented using `itertools.product`
for improved performance and memory efficiency.

✳️ Replace X.Y.Z with the release version when merging.

Related Issue

Closes: #28921

Author Note

👋 I'm a first-time contributor to NumPy and would love feedback on this PR. Thanks to @seberg and @ricardoV94 for the helpful context in the original issue!

📎 GitHub: @samruddhibaviskar11

tylerjereddy · 2025-06-10T23:19:03Z

benchmarks/benchmarks/bench_ndindex.py

+    print(f"Time to consume first 1000 elements: {partial_time * 1000:.3f} ms")
+    print("--> Memory efficient: Only generates indices as needed")
+
+if __name__ == "__main__":


It might be more sensible to follow our usual approach to writing/running benchmarks with asv. It has builtin harnesses for timing (time_) and peak memory usage (peakmem_). It will probably be easier to evaluate changes when sticking to this common approach we use for benchmarking perf/memory usage. It shouldn't be too hard to navigate around the other asv benchmarking modules we have to get a sort of "template" for how that works.

tylerjereddy · 2025-06-10T23:21:09Z

numpy/lib/_index_tricks_impl.py

+    New in version X.Y.Z.
+    As of NumPy 2.3.0.dev0, this iterator is implemented using `itertools.product`
+    from Python's standard library. This change provides significant improvements
+    in both performance and and memory efficiency, particularly for large iteration


repeated and; we'll probably want to change the version numbering you use above since i.e., 2.3.0 has already been released

tylerjereddy · 2025-06-10T23:25:45Z

numpy/lib/tests/test_index_tricks.py

+# 1.  Empty Shapes and Zero Dimensions
+def test_ndindex_empty_and_zero_dimensions():
+    # Empty shape
+    assert list(np.ndindex()) == [()]


It probably does make sense to add a few more test cases, but note that the already-existing test_ndindex already has this particular test case for example, and may cover one or two other cases.

tylerjereddy · 2025-06-10T23:27:08Z

numpy/lib/tests/test_index_tricks.py

+    with pytest.raises(TypeError):
+        list(np.ndindex([2, 3]))
+    with pytest.raises(TypeError):
+        list(np.ndindex((2.0, 3)))


minor: could use pytest.mark.parametrize for concision and separation of case execution

numpy/lib/_index_tricks_impl.py

tylerjereddy · 2025-06-10T23:30:01Z

numpy/lib/tests/test_index_tricks.py

+    expected_count = 20 * 30 * 40
+    assert_equal(count, expected_count)
+
+    assert elapsed_time < 1.0, f"ndindex took {elapsed_time:.3f}s, which seems slow"


we don't usually do this in our testsuite--we typically guard against performance regressions using asv benchmarks separate from our regression tests proper (see my other comment about asv)

eendebakpt · 2025-06-14T22:38:10Z

numpy/lib/_index_tricks_impl.py

+    
+    Notes
+    -----
+    New in version X.Y.Z.


This text seems fine for a release note. But since the interface is unchanged, I would leave it out if the user documentation.

jorenham

You should be able to fix the linter errors by running ruff check --fix locally. If you don't have ruff installed yet, then you can run pip install -r requirements/linter_requirements.txt to install the same ruff version as we use in CI.

numpy/lib/_index_tricks_impl.py

samruddhibaviskar11 · 2025-06-18T19:59:55Z

Hi @jorenham @tylerjereddy and @eendebakpt,

I've pushed updates that include:

Switched to product(*map(range, shape)) for cleaner ndindex initialization.

Ran ruff check --fix across all relevant files — no issues reported locally.

Amended the commit message and force-pushed.

That said, the CI lint jobs (ComprehensiveTests Lint, Linux tests / lint) are still failing — I’m currently investigating. Most other CI failures (macOS, Windows, BLAS, FreeBSD, Emscripten, etc.) appear unrelated to this PR.

The core focus remains the ndindex rewrite using itertools.product for better performance and memory usage, along with an expanded and refactored test suite.

Let me know if you'd like me to dive into any specific check. Thanks again for your feedback and support!

mattip

LGTM, needs a release note. Some of the comments in the tests are redundant

numpy/lib/_index_tricks_impl.py

numpy/lib/tests/test_index_tricks.py

eendebakpt · 2025-07-23T07:01:43Z

numpy/lib/_index_tricks_impl.py

-                       strides=_nx.zeros_like(shape))
-        self._it = _nx.nditer(x, flags=['multi_index', 'zerosize_ok'],
-                              order='C')
+        if any(dim < 0 for dim in shape):


Suggested change

if any(dim < 0 for dim in shape):

if min(shape) < 0:

Is perhaps faster, as it avoids creation of a generator.

eendebakpt · 2025-07-23T07:04:16Z

numpy/lib/tests/test_index_tricks.py

+def test_ndindex_negative_dimensions(negative_shape_arg):
+    """Test that negative dimensions raise ValueError."""
+    with pytest.raises(ValueError):
+        # ndindex should raise ValueError immediately for negative dimensions


If you are testing ndindex raises immediately, then remove the list in the next line

eendebakpt · 2025-07-23T07:04:28Z

numpy/lib/tests/test_index_tricks.py

+    """Test that negative dimensions raise ValueError."""
+    with pytest.raises(ValueError):
+        # ndindex should raise ValueError immediately for negative dimensions
+        list(ndindex(negative_shape_arg))


Suggested change

list(ndindex(negative_shape_arg))

ndindex(negative_shape_arg)

numpy/lib/__init__.py

eendebakpt · 2025-08-11T09:21:10Z

@samruddhibaviskar11 The PR looks like a nice improvement overall. Would you be able to address the remaining comments and look at the CI failures?

samruddhibaviskar11 · 2025-08-11T11:54:56Z

@samruddhibaviskar11 The PR looks like a nice improvement overall. Would you be able to address the remaining comments and look at the CI failures?

Yes, @eendebakpt I will look into the remaining comments and CI failures.

numpy/lib/__init__.py

…ficiency

…xpanded tests

…ts, and benchmarks

jorenham · 2025-08-11T22:40:16Z

Note that there are still several problems that are causing the CI to fail.

eendebakpt · 2025-08-12T08:01:23Z

numpy/lib/_index_tricks_impl.py

-                       strides=_nx.zeros_like(shape))
-        self._it = _nx.nditer(x, flags=['multi_index', 'zerosize_ok'],
-                              order='C')
+        if min(shape) < 0:


Suggested change

if min(shape) < 0:

if min(shape, default=0) < 0:

(to handle the case where shape is empty)

- Use min(..., default=0) to handle empty shapes safely. - Fix cascade issues in ndindex implementation. - Add regression tests to cover empty shapes and previous bugs.

numpy/lib/_index_tricks_impl.py

samruddhibaviskar11 · 2025-08-27T14:19:43Z

Hi all the CI is green ✅. Please take another look and let me know if anything else is needed. Thanks!

numpy/lib/_index_tricks_impl.py

doc/release/upcoming_changes/29165.performance.rst

eendebakpt · 2025-09-02T20:51:48Z

doc/release/upcoming_changes/29165.performance.rst

 providing significant improvements in both performance and memory efficiency,
 particularly for large iteration spaces, while maintaining the original behavior and interface.
-
+For example, for a (50, 60, 90) array, this results in a 5.2× speedup on average, 


Suggested change

For example, for a (50, 60, 90) array, this results in a 5.2× speedup on average,

For example, for an array of shape (50, 60, 90) the NumPy `npindex` benchmark improves performance by a factor 5.2.

Maybe leave of the memory. Is there really a significant amount saved compared to the size of the array?

Thanks! Updated to your phrasing and removed the memory mention as suggested. ✅

eendebakpt

I left a minor comment for the release notes, but otherwise this looks good.

Thanks @samruddhibaviskar11

jorenham

The code itself looks good now AFAIK. Once the other comments have also been addressed, and the other maintainers are also happy, this should be ready to merge.

samruddhibaviskar11 · 2025-09-05T18:18:15Z

Hi all, The CI doc build failed due to network issues fetching external references (pandas/matplotlib). Core tests are passing, so this seems like an infrastructure issue. Could someone help resolve this CI issue?

eendebakpt · 2025-09-05T18:26:42Z

Indeed looks like a network issue. You could either trigger a new CI build to be sure (make an empty commit), or the numpy maintainers could (if pr is approved) merge anyway.

samruddhibaviskar11 · 2025-09-05T18:29:11Z

Thanks for confirming! I’ll trigger a new CI build with an empty commit to be safe.

seberg

I think plenty reviewed and a quick look through, this looks nice, thanks a lot. Covered a lot of ground iwth benchmarks and tests.

samruddhibaviskar11 · 2025-09-06T08:39:39Z

Thank you all for reviewing and guiding me ✅ — I’ve learned a lot!

…29165) * PERF: Rewrite ndindex using itertools.product for speed and memory efficiency * Refactor ndindex using itertools.product and add ASV benchmarks and expanded tests * Fix linting issues flagged by Ruff across ndindex implementation, tests, and benchmarks * Add release note and finalize ndindex improvements * BUG: Raise ValueError for negative dimensions in ndindex * MAINT: Expose ndindex in numpy.lib for public usage and benchmarks * MAINT: Fix ndindex __module__ attribute for public API test * Remove numpy.distutils tests after dropping distutils in NumPy 2.0 * Format imports in numpy.lib.__init__.py with Ruff * Set ndindex.__module__ = 'numpy.lib' to match test expectation * BUG: Raise ValueError for negative dimensions in ndindex; MAINT: Remove ndindex from numpy.lib public API; TEST: Update negative dimension test * Revert distutils test deletions as requested * Revert numpy/lib/__init__.py to upstream/main state as requested * lib: ndindex—handle empty shape and fix implementation bugs - Use min(..., default=0) to handle empty shapes safely. - Fix cascade issues in ndindex implementation. - Add regression tests to cover empty shapes and previous bugs. * TST: remove imports to fix lint for test_index_tricks.py * STYLE: Remove ruff formatting and lint fixes * REF: ndindex test cases — added functional tests and applied Ruff fixes * Add benchmark comparing ndindex vs itertools.product * MAINT: remove unused in ndindex * DOC: clarify release note for ndindex performance improvement * trigger CI rebuild

samruddhibaviskar11 changed the title ~~PERF: Use itertools.product for ndindex to improve performance~~ ENH: Use itertools.product for ndindex to improve performance Jun 10, 2025

tylerjereddy added the 01 - Enhancement label Jun 10, 2025

tylerjereddy reviewed Jun 10, 2025

View reviewed changes

jorenham reviewed Jun 10, 2025

View reviewed changes

numpy/lib/_index_tricks_impl.py Outdated Show resolved Hide resolved

tylerjereddy reviewed Jun 10, 2025

View reviewed changes

eendebakpt reviewed Jun 14, 2025

View reviewed changes

jorenham reviewed Jun 15, 2025

View reviewed changes

numpy/lib/_index_tricks_impl.py Outdated Show resolved Hide resolved

melissawm added this to NumPy first-time contributor PRs Jun 19, 2025

melissawm moved this to Awaiting a code review in NumPy first-time contributor PRs Jun 19, 2025

mattip reviewed Jun 30, 2025

View reviewed changes

numpy/lib/_index_tricks_impl.py Outdated Show resolved Hide resolved

numpy/lib/tests/test_index_tricks.py Outdated Show resolved Hide resolved

samruddhibaviskar11 force-pushed the perf/ndindex-itertools branch from c75e95a to 94f2a73 Compare July 12, 2025 22:45

eendebakpt reviewed Jul 23, 2025

View reviewed changes

eendebakpt reviewed Aug 11, 2025

View reviewed changes

numpy/lib/__init__.py Outdated Show resolved Hide resolved

jorenham reviewed Aug 11, 2025

View reviewed changes

numpy/lib/__init__.py Show resolved Hide resolved

This comment was marked as resolved.

Sign in to view

samruddhibaviskar11 added 7 commits August 11, 2025 20:09

PERF: Rewrite ndindex using itertools.product for speed and memory ef…

fe0ea11

…ficiency

Refactor ndindex using itertools.product and add ASV benchmarks and e…

aac39b0

…xpanded tests

Fix linting issues flagged by Ruff across ndindex implementation, tes…

c558680

…ts, and benchmarks

Add release note and finalize ndindex improvements

481beaf

BUG: Raise ValueError for negative dimensions in ndindex

db99b82

MAINT: Expose ndindex in numpy.lib for public usage and benchmarks

e4295e3

MAINT: Fix ndindex __module__ attribute for public API test

1f42df7

samruddhibaviskar11 requested a review from jorenham August 11, 2025 21:09

eendebakpt reviewed Aug 12, 2025

View reviewed changes

lib: ndindex—handle empty shape and fix implementation bugs

cf267eb

- Use min(..., default=0) to handle empty shapes safely. - Fix cascade issues in ndindex implementation. - Add regression tests to cover empty shapes and previous bugs.

samruddhibaviskar11 force-pushed the perf/ndindex-itertools branch from a4cb71d to cf267eb Compare August 16, 2025 14:02

TST: remove imports to fix lint for test_index_tricks.py

fb11a4d

jorenham reviewed Aug 16, 2025

View reviewed changes

numpy/lib/_index_tricks_impl.py Show resolved Hide resolved

samruddhibaviskar11 added 3 commits August 19, 2025 17:28

STYLE: Remove ruff formatting and lint fixes

b87ae25

REF: ndindex test cases — added functional tests and applied Ruff fixes

549a5ca

Add benchmark comparing ndindex vs itertools.product

9079cb5

jorenham self-requested a review August 27, 2025 14:58

jorenham requested changes Aug 27, 2025

View reviewed changes

numpy/lib/_index_tricks_impl.py Outdated Show resolved Hide resolved

doc/release/upcoming_changes/29165.performance.rst Outdated Show resolved Hide resolved

MAINT: remove unused in ndindex

b670436

eendebakpt reviewed Sep 2, 2025

View reviewed changes

eendebakpt approved these changes Sep 2, 2025

View reviewed changes

jorenham self-requested a review September 3, 2025 17:55

jorenham approved these changes Sep 3, 2025

View reviewed changes

DOC: clarify release note for ndindex performance improvement

f2bbc98

trigger CI rebuild

a55f238

seberg approved these changes Sep 5, 2025

View reviewed changes

seberg merged commit 1a9d731 into numpy:main Sep 5, 2025
76 checks passed

github-project-automation bot moved this from Awaiting a code review to Completed in NumPy first-time contributor PRs Sep 5, 2025

ericpre mentioned this pull request Dec 5, 2025

Fix flyiter for numpy 2.4.0rc1 hyperspy/hyperspy#3571

Merged

4 tasks

	list(ndindex(negative_shape_arg))
	ndindex(negative_shape_arg)

	For example, for a (50, 60, 90) array, this results in a 5.2× speedup on average,
	For example, for an array of shape (50, 60, 90) the NumPy `npindex` benchmark improves performance by a factor 5.2.

Uh oh!

ENH: Use itertools.product for ndindex to improve performance #29165

ENH: Use itertools.product for ndindex to improve performance #29165

Uh oh!

Conversation

samruddhibaviskar11 commented Jun 10, 2025

Summary

Motivation

Changes Made

Benchmark Results

Memory Usage Test:

Test Cases & Functional Correctness

Environment/Testing Notes:

Docstring Update

Related Issue

Author Note

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

tylerjereddy Jun 10, 2025 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

jorenham left a comment

Choose a reason for hiding this comment

Uh oh!

Uh oh!

samruddhibaviskar11 commented Jun 18, 2025

Uh oh!

mattip left a comment

Choose a reason for hiding this comment

Uh oh!

Uh oh!

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Uh oh!

eendebakpt commented Aug 11, 2025

Uh oh!

samruddhibaviskar11 commented Aug 11, 2025

Uh oh!

Uh oh!

This comment was marked as resolved.

Uh oh!

jorenham commented Aug 11, 2025

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Uh oh!

samruddhibaviskar11 commented Aug 27, 2025

Uh oh!

Uh oh!

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

eendebakpt left a comment

Choose a reason for hiding this comment

Uh oh!

jorenham left a comment

Choose a reason for hiding this comment

Uh oh!

samruddhibaviskar11 commented Sep 5, 2025 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Uh oh!

eendebakpt commented Sep 5, 2025

tylerjereddy Jun 10, 2025 •

edited

Loading

samruddhibaviskar11 commented Sep 5, 2025 •

edited

Loading