Feature: NumPy-compliant distributed advanced indexing

[ClaudiaComito]

Description

TL;DR

This PR replaces Heat's legacy, local-only indexing with 100% NumPy-API-compliant advanced indexing capabilities across distributed nodes.

You can now seamlessly use boolean masks, integer arrays, negative-step slices etc. on distributed arrays without manual data shuffling:

import heat as ht

# Load a massive distributed dataset (e.g., 100 million rows, split across nodes/GPUs)
data = ht.random.randn(100_000_000, 50, split=0) 

#  distributed filtering condition 
outliers_mask = (data > 3.0).any(axis=1)

# extract only the outliers without gathering to a single node
outliers = data[outliers_mask]

---

This pull request introduces a significant overhaul of distributed indexing within dndarray.py, specifically targeting the __getitem__ and __setitem__ methods.

The logic has been completely refactored to identify zero-communication paths ("early out") for standard slices, while routing heavy, unordered (non-sequential) advanced indexing through highly optimized MPI collective communication.

Also, indexing.nonzero(), the kwarg as_tuple has been introduced (default: True) to comply with the Numpy API while giving users the choice to switch to torch-style output (2-D array).

Main changes (LAST UPDATE 13.6.2026)

dndarray.py

• replaced scattered key parsing with a centralized _resolve_indexing_state helper. This function torch-proofs all key inputs, handles broadcasting, aligns array dimension to indexed shape, and determines the indexing operation type for later dispatching. Returns a structured ProcessedKey NamedTuple.
• refactored the monolithic __getitem__ and __setitem__ functions. They are now wrappers that call the resolution state and dispatch to dedicated methods (e.g., __getitem_scalar, __setitem_mask, __getitem_advanced_local).
• added full support for non-sequential distributed advanced indexing (using integer arrays or boolean masks). It uses MPI.Alltoallv for cross-rank data fetching and assignment (__getitem_unordered and __setitem_unordered).
• introduced _resolve_duplicate_indices to guarantee NumPy-compliant "last assignment wins" semantics when using advanced indexing with duplicate indices on GPUs (thanks @Hakdag97 ).
• implemented support for descending/negative-step slices across distributed memory __getitem_descending_slice_distributed and its setter counterpart).
• added a __broadcast_value helper to automatically broadcast assigned values to perfectly match the target slice or boolean mask shape during __setitem__ operations.
• updated __torch_proxy__ to explicitly track the split axis natively within the tensor's named dimensions for safer split axis tracking during dimensions-changing operations.
• refactored tests, removed orphaned legacy code and tests.
• introduced INDEXING.md in doc/source/ and added it to the .rst index

indexing.py

• *modified nonzero to return a tuple of 1D DNDarrays by default (one array per dimension) instead of a single 2D coordinate matrix (Numpy API compliance).
• added as_tuple argument to nonzero to allow toggling between the new NumPy-style tuple output (True) and the legacy Torch-style 2D array output (False).
• Updated where(cond) to rely on nonzero and return consistent output (tuple of 1-D arrays) independently of split axis
• updated tests and docstrings

The following table is from doc/source/INDEXING.md

Array is distributed	Operation	Key is distributed	Value is distributed	Result is distributed	Notes
No	array[key]	No	--	No	Standard local indexing.
No	array[key]	Yes	--	Yes	The resulting array inherits the split axis and balanced status directly from the distributed key.
Yes	array[key]	No	--	Yes / No	No if the key is a pure scalar along the split axis (the split dimension is lost and the result is broadcasted). Yes for slices/masks. Non-sequential local advanced indices are automatically distributed across the split axis under the hood.
Yes	array[key]	Yes	--	Yes	Split axis is retained or shifted. Evaluated as a distr_mask fast-path or triggers __getitem_unordered for cross-node MPI collective fetching.
No	array[key] = val	No	No	No (In-place)	Standard local assignment.
Yes	array[key] = val	No	No	Yes (In-place)	The local value is automatically converted into a distributed array and broadcasted to align with the array's distribution constraints.
Yes	array[key] = val	No	Yes	Yes (In-place)	Split axis match required: If the value's split axis doesn't match the target's split axis, a RuntimeError is raised. If they do match, value is dynamically load-balanced (redistribute_) to match the target's chunk sizes before assignment.
Yes	array[key] = val	Yes	No, scalar	Yes (In-place)	A pure scalar value is correctly assigned to all masked/indexed elements across all MPI ranks natively.
Yes	array[key] = val	Yes	No, array	ERROR / Yes	Exception raised for integer indices. Supported for boolean masks via MPI prefix sums to dynamically slice the non-distributed array.
Yes	array[key] = val	Yes	Yes	Yes (In-place)	Communication-heavy: For masks, value is redistributed to match key. For integer arrays, key is redistributed to match value. Both are followed by an Alltoallv shuffle.

Note: Extracting a single element along the split axis will collapse that dimension, resulting in split=None.

Internal getitem/setitem routing logic

LAST UPDATE 13.6.2026

graph TD
    Start((Receive Key)) --> CheckScalar{Is key a pure scalar<br/>and not boolean?}
    
    CheckScalar -- Yes --> EvalRoot{Compute root}
    EvalRoot --> OpScalar[op_type = 'scalar']
    
    CheckScalar -- No --> CheckFastPath{Matches distr_mask<br/>fast path?}
    
    CheckFastPath -- Yes & not tuple --> OpDistrMask1[op_type = 'distr_mask']
    
    CheckFastPath -- No / Tuple --> Normalize[Normalize keys, extract bounds,<br/>check dimensionality & broadcast]
    
    Normalize --> FinalRouting{Evaluate Key State}
    
    FinalRouting -->|root is not None| OpScalar2[op_type = 'scalar']
    FinalRouting -->|split_key_is_ordered == 0| OpDist[op_type = 'distributed'<br/>Unordered MPI Communication]
    FinalRouting -->|split_key_is_ordered == -1| OpDesc[op_type = 'descending_slice']
    
    FinalRouting -->|key_is_mask_like == True| MaskTypeCheck{distr_mask_fast_path?}
    MaskTypeCheck -- Yes --> OpDistrMask2[op_type = 'distr_mask']
    MaskTypeCheck -- No --> OpLocalMask[op_type = 'local_mask']
    
    FinalRouting -->|Default / Pure Slices / Ordered| OpAdv[op_type = 'advanced'<br/>Local Fast Path]

    %% Map to actual handlers
    subgraph Handlers [Target Routing Methods]
        OpScalar & OpScalar2 --> H_Scalar[__getitem_scalar<br/>__setitem_scalar]
        OpDist --> H_Dist[__getitem_advanced_distributed<br/>__setitem_advanced_distributed]
        OpDesc --> H_Desc[__getitem_descending_slice_distributed<br/>__setitem_descending_slice_distributed]
        OpDistrMask1 & OpDistrMask2 --> H_DistMask[__getitem_mask<br/>__setitem_mask]
        OpLocalMask --> H_LocalMask[__getitem_advanced_local<br/>__setitem_advanced_local]
        OpAdv --> H_Adv[__getitem_advanced_local<br/>__setitem_advanced_local]
    end
    
    %% Styling
    classDef target fill:#d4edda,stroke:#28a745,stroke-width:2px;
    class H_Scalar,H_Dist,H_Desc,H_DistMask,H_LocalMask,H_Adv target;

Loading

Memory footprint

Scaling behaviour

Issue/s resolved: #703 #914 #918 #1012 #1019 #2135 #1816 #824

Type of change

• Breaking change (nonzero() is now Numpy-compliant by default and returns a tuple of 1-D arrays)

Memory requirements

Performance

Due Diligence

• All split configurations tested
• Multiple dtypes tested in relevant functions
• Documentation updated (if needed)
• Updated changelog.md under the title "Pending Additions"

Does this change modify the behaviour of other functions? If so, which?

yes, everything that relied on the legacy indexing quirks (fixed) and everything that relied on 2D output from nonzero() (also fixed)

[ghost]

👇 Click on the image for a new way to code review

Legend

[brownbaerchen]

I didn't manage to read everything in this review pass. I think there is still plenty to do, mainly in terms of casting to DNDarray in the beginning in order to avoid excessive checks for various cases, and isolating some cases into their own functions.

[brownbaerchen]

Suggested change

	try:
	# is key an ndarray or DNDarray or torch.Tensor?
	key = key.item()
	except AttributeError:
	# key is already an integer, do nothing
	pass
	# cast key to scalar if it is an array
	try:
	key = key.item()
	except AttributeError:
	pass

[brownbaerchen]

This is not consistent with the type hint that root is integer

[brownbaerchen]

Again, not consistent with type hints.

[brownbaerchen]

Please give some information about what the return values of this function are.

[brownbaerchen]

Can you provide some examples of what this is doing? It's not clear to me from the title and description alone what this function does and why.

[brownbaerchen]

This could be much simpler if we knew what kind of array k is if k is array like

[brownbaerchen]

Suggested change

	if ellipsis > 1:
	raise ValueError("indexing key can only contain 1 Ellipsis (...)")
	if ellipsis:
	# key contains exactly 1 ellipsis
	if ellipsis > 1:
	raise ValueError("indexing key can only contain 1 Ellipsis (...)")
	elif ellipsis == 1:

[brownbaerchen]

This can be simplified if we knew what type k is by previous casting.

[brownbaerchen]

Is this while loop needed? Would it be possible to expand all dims at once and setup all slices at once?

[brownbaerchen]

what is lose_dims?

[JuanPedroGHM]

There are a lot of comments from @brownbaerchen that it is unclear if they have been addressed. Other than that, I would say this is ready to go.

[JuanPedroGHM]

Documentation looks great!