Skip to content

[core][ndarray] Fix shape reconstruction logic in NDArray.__getitem__()#326

Merged
forfudan merged 2 commits into
Mojo-Numerics-and-Algorithms-group:pre-0.9from
josiahls:fix/ndarray-int-variadic-type-list-mismatch
Mar 1, 2026
Merged

[core][ndarray] Fix shape reconstruction logic in NDArray.__getitem__()#326
forfudan merged 2 commits into
Mojo-Numerics-and-Algorithms-group:pre-0.9from
josiahls:fix/ndarray-int-variadic-type-list-mismatch

Conversation

@josiahls
Copy link
Copy Markdown
Contributor

@josiahls josiahls commented Mar 1, 2026

Caused by:

ValueError: Number of dimensions must be positive, got 0. Provide at least one shape dimension. [at NDArrayShape.__init__(shape: List[Int])]

on

print(array[0, 0])

Additionally added a TODO since the error message was otherwise extremely vague.


Updates during review by @forfudan:

index_type_list.append(IndexTypeInfo(is_integer=True)), without which the mixture of slices and integers would cause integers to be missed in the record.

However, this PR is not sufficient to address the problem. There are some other bugs in the code related to shape reconstruction logic, which will be even exposued after adding the index_type_list.

Here are some further fixes:

  1. Fixed is_integer branch: The old code new_shape.append(narr.shape[i]) was wrong in two ways — it used i (position in index_type_list) instead of a proper dimension counter, and it added the dimension to the output shape when integer indexing should remove it (dimension reduction). Changed to just advance narr_dim without appending (skip it).

  2. Fixed missing trailing dimensions: When index_type_list has fewer entries than narr's dimensions (e.g., b[Slice(0,2)] on a 2D array), remaining dimensions were silently dropped. Added a while narr_dim < ndims loop to preserve them.

  3. Fixed ellipsis dimension counting: Changed from self.ndim - len(index_type_list) + 1 (which incorrectly assumes every non-ellipsis entry consumes a dimension — NewAxis doesn't) to counting only is_slice and is_integer entries.

  4. Removed stale TODO from the original PR: # TODO: Shouldn't we assert len(shape_size) == len(slice_list)? referenced a nonexistent variable shape_size.

@josiahls
Copy link
Copy Markdown
Contributor Author

josiahls commented Mar 1, 2026

Note it looks like a lot of unit tests in NuMojo/tests/core/test_array_indexing_and_slicing.mojo are commented out, which might explain why this got through.

Caused by:
```
ValueError: Number of dimensions must be positive, got 0. Provide at least one shape dimension. [at NDArrayShape.__init__(shape: List[Int])]
```
on
```
print(array[0, 0])
```

Additionally added a TODO since the error message was otherwise extremely vague.
@josiahls josiahls force-pushed the fix/ndarray-int-variadic-type-list-mismatch branch from 5d0d2ad to 4ed56dc Compare March 1, 2026 15:04
@forfudan forfudan self-requested a review March 1, 2026 16:11
@forfudan forfudan added the bug Something isn't working label Mar 1, 2026
@forfudan forfudan changed the title [core][ndarray]: Fix: Missing index_type_list in NDArray.__getitem__ [core][ndarray] Fix shape reconstruction logic in NDArray.__getitem__() Mar 1, 2026
@forfudan
Copy link
Copy Markdown
Collaborator

forfudan commented Mar 1, 2026

@josiahls, thank you for the PR.

Your PR catched a missing index_type_list.append(IndexTypeInfo(is_integer=True)), without which the mixture of slices and integers would cause integers to be missed in the record.

However, this PR is not sufficient to address the problem. There are some other bugs in the code related to shape reconstruction logic, which will be even exposued after adding the index_type_list.

So I made some other fixes on top of your PR, changed upon yours, and update the description.


@shivasankarka, I am adding tests for this case, please let me handle the approval for this PR. Thanks :D

Copy link
Copy Markdown
Collaborator

@forfudan forfudan left a comment

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Approved with some further fixes, need to also add some tests in another PR.

@forfudan forfudan merged commit 842a67d into Mojo-Numerics-and-Algorithms-group:pre-0.9 Mar 1, 2026
2 checks passed
shivasankarka pushed a commit that referenced this pull request Mar 2, 2026
… refine them, and add some more (#327)

This PR is related to #326. It uncomments tests for indexing and slicing
of NDArray, refines them, and some more tests.
@shivasankarka
Copy link
Copy Markdown
Collaborator

@josiahls @forfudan Thanks for the fix :)

forfudan pushed a commit that referenced this pull request Mar 31, 2026
…_()` (#326)

This PR fixes shape reconstruction logic in `NDArrray.__getitem__()` dunder method.

1. Added `index_type_list.append(IndexTypeInfo(is_integer=True))`, without which the mixture of slices and integers would cause integers to be missed in the record.

2. Fixed `is_integer` branch: The old code `new_shape.append(narr.shape[i])` was wrong in two ways — it used `i` (position in `index_type_list`) instead of a proper dimension counter, and it added the dimension to the output shape when integer indexing should remove it (dimension reduction). Changed to just advance `narr_dim` without appending (skip it).

3. Fixed missing trailing dimensions: When `index_type_list` has fewer entries than narr's dimensions (e.g., `b[Slice(0,2)]` on a 2D array), remaining dimensions were silently dropped. Added a while `narr_dim < ndims` loop to preserve them.

4. Fixed ellipsis dimension counting: Changed from `self.ndim - len(index_type_list) + 1` (which incorrectly assumes every non-ellipsis entry consumes a dimension — NewAxis doesn't) to counting only `is_slice` and `is_integer` entries.
forfudan added a commit that referenced this pull request Mar 31, 2026
… refine them, and add some more (#327)

This PR is related to #326. It uncomments tests for indexing and slicing
of NDArray, refines them, and some more tests.
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

bug Something isn't working

Projects

None yet

Development

Successfully merging this pull request may close these issues.

3 participants