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

[mrfh92]

just a comment: in the fft-module (if already merged at time merging this PR) some commented-out parts of test_rfftn_irfftn can be added again if this PR is ready

[JuanPedroGHM]

Why was the ellipsis removed? This looks incorrect to me.

[JuanPedroGHM]

This never returns "slice", so I think I think those branches of getitem and setitem are never reached.

[JuanPedroGHM]

I am not sure how relevant this is, but I believe this should also check for boolean arrays wrapped in a tuple. In numpy, a[(mask,)] is the same as a[mask]. It is weird non-documented behaviour in numpy, as far as I can tell.

Here is an example that reproduces the error in any number of ranks, as far as I can tell.

import numpy as np
import heat as ht
from mpi4py import MPI

# Numpy
a_np = np.arange(60, dtype=np.float32).reshape(5, 4, 3)
mask_np = a_np > 30
if MPI.COMM_WORLD.Get_rank() == 0:
    print("NUMPY")
    print(a_np[mask_np])
    print(a_np[(mask_np,)])

# Heat
    print("HEAT")

a = ht.array(a_np, split=0)
print(a)
mask = a > 30
print(mask)

print(a[mask])
print(a[(mask,)])

[brownbaerchen]

Suggested change

	class LocalIndex:
	"""
	Indexing class for local operations (primarily for :func:`lloc` function)
	For docs on ``__getitem__`` and ``__setitem__`` see :func:`lloc`
	"""
	def __init__(self, obj):
	self.obj = obj
	def __getitem__(self, key):
	return self.obj[key]
	def __setitem__(self, key, value):
	self.obj[key] = value

this was removed in #2494