diff --git a/doc/source/INDEXING.md b/doc/source/INDEXING.md
new file mode 100644
index 0000000000..bcd8005740
--- /dev/null
+++ b/doc/source/INDEXING.md
@@ -0,0 +1,258 @@
+
+# Indexing on DNDarrays
+
+Heat v2.0 introduces fully distributed indexing for DNDarrays. While the indexing behaviour is designed to be highly compatible with the NumPy API, the memory-distributed nature of DNDarrays introduces unique considerations regarding performance and communication overhead. In the following sections, we will cover the basics plus some of these Heat-specific indexing features.
+
+*Note: This guide is heavily inspired by the official [NumPy indexing documentation](https://numpy.org/doc/stable/user/basics.indexing.html).*
+
+## Distributed indexing
+
+We work under the assumption that Heat users process data in very large, memory-distributed arrays. In the following, we will refer to `array`, `key`, and `value` as the DNDarray, the index/combination of indices, and (if present) the value to be assigned to the index, respectively. Examples:
+
+- item getting: `array[key]`
+- item setting: `array[key] = value`
+
+We assume that not only `array`, but also `key` and `value` may be very large and distributed across MPI processes if the use case requires.
+
+The following table shows the distribution semantics of the DNDarray indexing operations.
+
+| 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`.*
+
+---
+
+## Basic slicing and indexing
+
+Basic slicing extends Python's basic concept of slicing to N dimensions. It occurs when `key` is a `slice` object (constructed by `start:stop:step` notation inside brackets), an integer, or a tuple of slice objects and integers.
+
+### Single element indexing
+When indexing a single element or a specific slice that reduces the dimensionality of the array, the `split` axis is dynamically updated. If the array is indexed with an integer along the dimension it is split on, that dimension is collapsed and the resulting slice is no longer distributed along that axis.
+
+```python
+import heat as ht
+
+# 1D array distributed across processes
+x = ht.arange(10, split=0)
+# indexing collapses the 0th dimension; the result is no longer distributed
+result = x[2]
+# result.split is None
+```
+
+If the array is multi-dimensional and split on an axis that is not the one being collapsed, the split axis shifts to account for the removed dimension.
+
+```python
+# 2D array distributed along axis 1 (columns)
+x = ht.arange(10).reshape(2, 5)
+x_split1 = ht.array(x, split=1)
+
+# selecting a specific row collapses axis 0
+result = x_split1[0]
+# result.split is 0, because the old axis 1 is now the new axis 0
+```
+
+### Slicing and striding
+Standard slicing `start:stop:step` preserves the dimensions of the array. The array remains distributed along the original split axis. Negative steps are supported and will reverse the elements locally while executing collective communication to reverse the chunks globally.
+
+```python
+x = ht.arange(20, split=0)
+# slice with a step
+result = x[1:11:3]
+# result.split remains 0
+```
+
+### Dimensional indexing
+
+You can manipulate the dimensionality of a DNDarray directly inside the brackets using ht.newaxis (or None) and ... (Ellipsis).
+
+- `None` or `np.newaxis` inserts a new axis of size 1 into the array's shape. If the array is distributed, inserting an axis before the split axis will cause the split axis index to shift by +1.
+
+- `...` expands to the number of `:` objects needed to make a selection tuple of the same length as the array dimensions.
+
+```python
+x = ht.array([[[1], [2], [3]], [[4], [5], [6]]], split=1)
+
+# adds a new dimension at axis 1
+x_newaxis = x[:, None, :2, :]
+# original split was 1; new split is 2
+```
+
+## Advanced indexing
+
+Advanced indexing is triggered when the selection object key is a non-tuple sequence object, a DNDarray (of integer or boolean data type), a torch.Tensor, or a tuple with at least one sequence object or multi-dimensional array.
+
+Advanced indexing always returns a copy of the data (contrast with basic slicing that returns a view).
+
+### Integer array indexing
+
+You can use DNDarray objects containing integers to select arbitrary items. The resulting array will take on the distribution map of the indexing key.
+
+
+```python
+# array split along axis 0
+x = ht.arange(60, split=0).reshape(5, 3, 4)
+
+# using multiple non-distributed indices
+k1 = ht.array([0, 4, 1, 0])
+k2 = ht.array([0, 2, 1, 0])
+k3 = ht.array([1, 2, 3, 1])
+
+# standard advanced indexing
+result = x[k1, k2, k3]
+```
+
+### Boolean array indexing
+
+Boolean arrays used as indices are treated as a mask. The result is a 1-D array containing the elements that correspond to True in the boolean array.
+
+```python
+arr = ht.arange(60, split=0).reshape(3, 4, 5)
+mask = arr > 30
+
+# returns a 1D array of all elements > 30, split along axis 0
+result = arr[mask]
+```
+
+Row-selection optimization: Heat implements a highly optimized fast-path for the common data science pattern of row-filtering. If you index a 2D array split along axis 0 with a 1D boolean mask that is also split along axis 0, Heat skips the heavy distributed indexing machinery. It applies the mask locally and resolves the global shape via a fast metadata exchange. The output remains a 2D array split along axis 0.
+
+
+```python
+arr_2d = ht.arange(20, split=0).reshape((10, 2))
+mask_1d = ht.array([True, False, True, False, True, False, True, False, True, False], split=0)
+
+# the result remains a 2D array (shape: 5, 2) and retains split=0
+result = arr_2d[mask_1d]
+```
+
+### In-place assignment (setitem)
+
+Advanced indexing can be used to assign values. If the assignment value is itself a distributed DNDarray, Heat will automatically execute a distributed routing protocol (via Alltoallv) to align the spatial memory distribution of the values with the target indices before executing the local assignments.
+
+```python
+x = ht.arange(10 * 20 * 30, split=1).reshape(10, 20, 30)
+
+# boolean mask assignment
+mask = x > 100
+x[mask] = 99.0
+
+# advanced integer assignment with a distributed value
+indices = ht.random.randint(0, 20, (2, 3, 4), dtype=ht.int64, split=0)
+value = ht.ones((1, 2, 3, 4, 1), split=0)
+
+# value is automatically broadcasted and redistributed to match 'x[..., indices, :]'
+x[..., indices, :] = value
+```
+
+## Combining advanced and basic indexing
+
+When you mix advanced indexing (like integer arrays or lists) with basic slicing (like `:`), the shape of the resulting `DNDarray` depends on whether the advanced indices are positioned next to each other.
+
+Heat follows NumPy's standard transposition rules for mixed indexing, while automatically managing the distributed memory alignment internally. The array's nominal `split` axis will track the new dimensional layout.
+
+### Advanced indexing on consecutive dimensions
+If the advanced indices are adjacent to each other (not separated by a slice), the resulting broadcasted shape of the advanced indices is inserted directly into the output shape at the position of the first advanced index.
+
+If the original array's `split` axis is untouched by the advanced indexing, it will simply shift to account for the collapsed dimensions.
+
+```python
+import heat as ht
+
+# arr shape: (10, 20, 30, 40), distributed along axis 3
+arr = ht.zeros((10, 20, 30, 40), split=3)
+a1 = ht.array([1, 2])
+a2 = ht.array([3, 4])
+
+# Advanced indices are consecutive on axes 1 and 2
+result = arr[:, a1, a2, :]
+
+# The advanced indices on axes 1 and 2 broadcast to a single shape (2,)
+# Result shape: (10, 2, 40)
+
+# The original split axis 3 is now the last dimension in the new shape.
+# result.split is 2
+```
+
+### Advanced indexing on non-consecutive dimensions
+
+If the advanced indices are separated by a basic slice, the resulting layout becomes ambiguous. To resolve this, the advanced-indexing dimensions are grouped together and transposed to the very front of the resulting array's shape.
+
+Any remaining basic slices follow behind them. The split axis is tracked through this transposition and assigned its new relative index.
+
+```python
+import heat as ht
+
+# arr shape: (10, 20, 30, 40), distributed along axis 3
+arr = ht.zeros((10, 20, 30, 40), split=3)
+a1 = ht.array([1, 2])
+a2 = ht.array([3, 4])
+
+# Advanced indices (axes 0 and 2) are separated by a slice (axis 1)
+result = arr[a1, :, a2, :]
+
+# The advanced indices broadcast to shape (2,) and are moved to the front.
+# The untouched basic slices (from axes 1 and 3) are appended to the back.
+# Result shape: (2, 20, 40)
+
+# The original split axis 3 is still the last dimension in the new array.
+# result.split is 2
+```
+
+## Communication overhead
+
+The indexing operations dynamically evaluate the state of the indexing key to determine the most efficient network routing strategy. The communication overhead ranges from completely zero (purely local execution) to heavy all-to-all exchanges for non-sequential advanced indexing.
+
+Here are the different possible configurations, categorized and ordered from the lowest communication overhead to the highest within each category.
+
+### Summary of Communication Overhead
+
+| Category | Configuration (Operation & State) | Communication Overhead (MPI Calls) |
+| :--- | :--- | :--- |
+| **Single Element Indexing** | `array[key]` (key is an integer on a *non-split* axis) | **None** |
+| | `array[key] = local_value` (key is an int on the *split* axis) | **None** (Only the root rank executes the local set) |
+| | `array[key]` (key is an int on the *split* axis) | **1 `Bcast`** (Root extracts value and broadcasts to all ranks) |
+| **Slicing & Striding** | `array[slice]` or `array[slice] = local_value` | **None** |
+| | `array[::-1]` (Descending slice along split axis) | **None** (Executes local slice followed by a global `flip` operation) |
+| | `array[::-1] = distributed_value` (Descending slice write) | **Multiple `Send`/`Recv`** (Executes `redistribute_` using point-to-point transfers if array slice and value are misaligned) |
+| **Dimensional Indexing** | `array[..., None]` or `array[:, np.newaxis]` | **None** |
+| **Advanced Indexing** | `array[mask] = local_value` (Boolean mask assignment) | **None / 1 `exscan`** (Zero for scalars; requires prefix sum for 1D local arrays) |
+| | `array[mask]` (1D bool mask on 2D array, both split=0) | **None** (Locally applied, delegates global shape resolution to `factories.array`) |
+| | `array[non_seq_key] = local_value` (Integer array assignment) | **2 `Allreduce`** (Evaluates global key bounds, then applies locally) |
+| | `array[non_seq_key]` (Standard unstructured advanced read) | **1 `Allgather` + 2 `Alltoallv`** (Builds comm matrix, requests indices, returns data) |
+| | `array[non_seq_key] = distributed_value` | **2 `Allreduce` + 1 `Allgather` + 1 `Alltoallv`** (+ hidden P2P in `redistribute_` if necessary) |
+
+---
+
+### Detailed Breakdown by Category
+
+#### 1. Slicing and striding
+* **Zero Overhead (Fast Path):** If the key consists solely of basic components (slices with positive steps, integers, `None`, or `...`), the `_resolve_indexing_state` method dynamically assigns an `op_type` (like `"slice"` or `"scalar"`) that bypasses state-checking `Allreduce` calls entirely, resulting in zero MPI overhead.
+* **Negative Slicing / Descending Strides (Low to Moderate Overhead):** Because PyTorch does not natively support negative slice steps, descending slices (e.g., `[::-1]`) are caught during key processing and explicitly converted into integer tensors (`torch.arange`).
+ * For **Reads** (`__getitem__`), the `op_type` is evaluated as `"descending_slice"`. This bypasses non-sequential routing and triggers `__getitem_descending_slice_distributed`, which performs a local slice and wraps the result in an unbalanced array before executing a global `flip` operation.
+ * For **Writes** (`__setitem__`), a specific handler for `"descending_slice"` flips the right-hand value and dynamically matches its distribution map to the key, triggering point-to-point `Send`/`Recv` exchanges via the `redistribute_` method.
+
+#### 2. Dimensional indexing
+* **Zero Overhead:** The use of `None`, `np.newaxis`, or `...` (Ellipsis) is handled during the initial `__process_key` phase. These simply manipulate the local array dimensions and update the split axis bookkeeping without requiring cross-rank data movement. Because they evaluate as basic components, they hit the zero-overhead Fast Path.
+
+#### 3. Single element indexing
+* **Zero Overhead (Non-Split Axis Get/Set):** If the scalar index is applied to any dimension other than the split dimension, the operation evaluates locally on all ranks.
+* **Zero Overhead (Split Axis Set):** In `__setitem__`, assigning a local or scalar value to a single index on the split axis identifies a `root` process. Only the `root` process performs the assignment; no broadcast is performed.
+* **Low Overhead (Split Axis Get):** In `__getitem__`, if a single element is requested along the split axis, the `root` process extracts the local tensor and uses a single `MPI.Bcast` to share the result with all other ranks.
+
+#### 4. Advanced indexing (integer array and boolean array)
+Advanced indexing covers the most complex routing logic, where overhead scales based on the operation and the nature of the value being assigned.
+* **Zero to low (local mask assignment):** When assigning a scalar using a boolean mask, the operation evaluates locally with zero MPI overhead. If assigning a 1D non-distributed tensor to an N-D mask, it uses a single `MPI.exscan` to compute sequence offsets.
+* **Very low (local integer assignment):** If a local array is assigned using integer arrays, the system performs two `MPI.Allreduce` calls to securely validate index bounds and check for negative indices across ranks. Afterwards, it isolates the assignment using `_advanced_setitem_unordered_local`, avoiding heavy payload exchanges.
+* **Low (row-selection optimization):** A dedicated fast path exists for 2D arrays split along axis 0 when indexed by a 1D boolean mask. It avoids full advanced indexing by applying the mask locally.
+* **High (non-sequential read):** Standard advanced reads rely on a collective MPI approach. Ranks first execute an `MPI.Allgather` to build a communication matrix. Active ranks then distribute their requested indices via `MPI.Alltoallv`, execute local lookups, and return the resulting data via a second `MPI.Alltoallv` exchange.
+* **Very high (distributed assignment):** If `__setitem__` is called with a distributed value, the engine must align the distributions using point-to-point communications (`redistribute_`), use an `Allgather` to construct the communication matrix, and shuffle the payload concurrently via an `MPI.Alltoallv` operation.
diff --git a/doc/source/_static/css/custom.css b/doc/source/_static/css/custom.css
index 553e7e5ea3..68c9b6f74c 100644
--- a/doc/source/_static/css/custom.css
+++ b/doc/source/_static/css/custom.css
@@ -7475,3 +7475,17 @@ span[id*='MathJax-Span'] {
padding: 20px;
border-top: 1px solid var(--hgrey-light);
}
+
+/* ============================================================
+ CUSTOM OVERRIDES: Force hyperlinks to be visible in italics
+ ============================================================ */
+.rst-content a,
+.md-typeset a {
+ color: #2980B9 !important; /* Matches standard RTD blue */
+ text-decoration: none;
+}
+
+.rst-content a:hover,
+.md-typeset a:hover {
+ text-decoration: underline !important;
+}
diff --git a/doc/source/index.rst b/doc/source/index.rst
index 6c52b4817c..cafdb51220 100644
--- a/doc/source/index.rst
+++ b/doc/source/index.rst
@@ -20,6 +20,7 @@ Release: |release|
documentation_howto
CONTRIBUTING
CODE_OF_CONDUCT
+ INDEXING
Also visit us on `GitHub `_ for more examples, docs, code and contributions.
diff --git a/heat/classification/kneighborsclassifier.py b/heat/classification/kneighborsclassifier.py
index 90d1859537..e438645f70 100644
--- a/heat/classification/kneighborsclassifier.py
+++ b/heat/classification/kneighborsclassifier.py
@@ -122,11 +122,11 @@ def predict(self, x: DNDarray) -> DNDarray:
"""
distances = self.effective_metric_(x, self.x)
_, indices = ht.topk(distances, self.n_neighbors, largest=False)
- predictions = self.y[indices.flatten()]
+
+ predictions = self.y[indices]
predictions.balance_()
- predictions = ht.reshape(predictions, (indices.gshape + (self.y.gshape[1],)))
+ predictions = ht.reshape(predictions, indices.gshape + (self.y.gshape[1],))
predictions = ht.sum(predictions, axis=1)
self.classes_ = ht.argmax(predictions, axis=1)
-
return self.classes_
diff --git a/heat/core/dndarray.py b/heat/core/dndarray.py
index 16ce355700..6272b05e0d 100644
--- a/heat/core/dndarray.py
+++ b/heat/core/dndarray.py
@@ -2,15 +2,13 @@
from __future__ import annotations
-import math
import numpy as np
import torch
import warnings
-from inspect import stack
from mpi4py import MPI
-from pathlib import Path
-from typing import List, Union, Tuple, TypeVar, Optional
+from typing import TypeVar, Any, Union
+from collections.abc import Iterable
warnings.simplefilter("always", ResourceWarning)
@@ -19,6 +17,10 @@
Communication = TypeVar("Communication")
+# Type aliases
+Index = Union[int, slice, type(...), None, torch.Tensor, np.ndarray, "DNDarray"]
+Indexer = Union[Index, tuple[Index, ...], list[Index]]
+
class LocalIndex:
"""
@@ -36,16 +38,926 @@ def __setitem__(self, key, value):
self.obj[key] = value
+from typing import NamedTuple
+
+
+class ProcessedKey(NamedTuple):
+ """
+ A named tuple to store the processed key information for distributed indexing operations.
+ """
+
+ key: Any
+ op_type: str # "scalar", "slice", "descending_slice", "distr_mask", "local_mask", "advanced", "distributed"
+ output_shape: tuple
+ output_split: int | None
+ split_key_is_ordered: int
+ key_is_mask_like: bool
+ out_is_balanced: bool
+ root: int | None
+ backwards_transpose_axes: tuple
+
+
+def _process_scalar_key(
+ arr: "DNDarray",
+ key: int | "DNDarray" | torch.Tensor | np.ndarray,
+ indexed_axis: int,
+ return_local_indices: bool | None = False,
+) -> tuple[int, int]:
+ """
+ Private helper function to process a single-item scalar key used for indexing a ``DNDarray``.
+ """
+ device = arr.larray.device
+ try:
+ # is key an ndarray or DNDarray or torch.Tensor?
+ key = key.item()
+ except AttributeError:
+ # key is already an integer, do nothing
+ pass
+ if not arr.is_distributed():
+ root = None
+ return key, root
+ if arr.split == indexed_axis:
+ # adjust negative key
+ if key < 0:
+ key += arr.gshape[indexed_axis]
+ # work out active process
+ _, displs = arr.counts_displs()
+ if key in displs:
+ root = displs.index(key)
+ else:
+ displs = torch.cat(
+ (
+ torch.tensor(displs, device=device),
+ torch.tensor(key, device=device).reshape(-1),
+ ),
+ dim=0,
+ )
+ _, sorted_indices = displs.unique(sorted=True, return_inverse=True)
+ root = sorted_indices[-1].item() - 1
+ displs = displs.tolist()
+ # correct key for rank-specific displacement
+ if return_local_indices:
+ if arr.comm.rank == root:
+ key -= displs[root]
+ else:
+ root = None
+ return key, root
+
+
+def _resolve_duplicate_indices(
+ key_in,
+ rhs_in: torch.Tensor,
+ target_shape: tuple[int, ...],
+):
+ """
+ CUDA-safe handling for duplicate advanced indices:
+ enforce NumPy semantics (last assignment wins) by dropping earlier duplicates.
+ Works for:
+ - key_in: torch.Tensor (indexes axis 0)
+ - key_in: tuple/list of torch.Tensors (pure advanced indexing)
+ rhs_in must match the indexing result shape.
+ """
+ # Scalars or single element: no need to deduplicate
+ if rhs_in.numel() <= 1:
+ return key_in, rhs_in
+
+ # Normalize key to either a single tensor or tuple of tensors
+ if torch.is_tensor(key_in):
+ idx_tensors = (key_in,)
+ elif (
+ isinstance(key_in, (tuple, list))
+ and len(key_in) > 0
+ and all(torch.is_tensor(k) for k in key_in)
+ ):
+ idx_tensors = tuple(key_in)
+ else:
+ # Not pure advanced-tensor indexing -> don't touch
+ return key_in, rhs_in
+
+ device = rhs_in.device
+
+ # Broadcast indices to common shape, then flatten
+ try:
+ idx_b = torch.broadcast_tensors(*idx_tensors)
+ except RuntimeError:
+ # If broadcast fails, leave it to PyTorch (will error appropriately)
+ return key_in, rhs_in
+
+ pos_shape = idx_b[0].shape
+ pos_ndim = len(pos_shape)
+ n = idx_b[0].numel()
+
+ idx_flat = [t.to(device=device, dtype=torch.int64).reshape(-1) for t in idx_b]
+
+ # Build linear index for duplicate detection
+ if len(idx_flat) == 1:
+ lin = idx_flat[0]
+ else:
+ lin = idx_flat[0]
+ # linearize across the first len(idx_flat) dimensions of the target tensor
+ for d in range(1, len(idx_flat)):
+ lin = lin * int(target_shape[d]) + idx_flat[d]
+
+ # Fast path: no duplicates
+ if torch.unique(lin).numel() == n:
+ return key_in, rhs_in
+
+ # Determine "last occurrence" per linear index (last wins)
+ pos = torch.arange(n, device=device, dtype=torch.int64)
+
+ # Prefer stable sort by lin if available; otherwise sort by combined key
+ try:
+ order = torch.argsort(lin, stable=True)
+ except TypeError:
+ # combined key sorts by lin, then by pos
+ combined = lin.to(torch.int64) * (n + 1) + pos
+ order = torch.argsort(combined)
+
+ lin_s = lin[order]
+ pos_s = pos[order]
+
+ is_last = torch.ones_like(lin_s, dtype=torch.bool)
+ is_last[:-1] = lin_s[1:] != lin_s[:-1]
+ keep_pos = pos_s[is_last] # positions in original stream
+
+ # Reduce RHS accordingly:
+ # Flatten leading "pos_ndim" dims into one, keep trailing dims as payload
+ rhs_view = rhs_in.reshape(n, *rhs_in.shape[pos_ndim:])
+ rhs_u = rhs_view[keep_pos].reshape(keep_pos.numel(), *rhs_in.shape[pos_ndim:])
+
+ # Reduce indices accordingly (use flattened 1D indices)
+ if torch.is_tensor(key_in):
+ key_u = idx_flat[0][keep_pos]
+ return key_u, rhs_u
+
+ key_u = tuple(t[keep_pos] for t in idx_flat)
+ return key_u, rhs_u
+
+
+def _resolve_indexing_state(
+ arr: "DNDarray",
+ key: Indexer,
+ return_local_indices: bool | None = False,
+ op: str | None = None,
+) -> tuple["DNDarray", ProcessedKey]:
+ """
+ Private helper function to align the indexing key and the array for distributed indexing operations.
+ This function is used internally by both ``__getitem__`` and ``__setitem__`` pipelines.
+
+ After processing the key, the following conditions are guaranteed:
+ - Any ellipses (`...`) or newaxis (`None`) objects have been replaced with the appropriate number of slice objects.
+ - ``np.ndarray`` and ``DNDarray`` objects have been converted to process-local ``torch.Tensor`` objects.
+ - The dimensionality of the key perfectly matches the (potentially modified) ``DNDarray`` it indexes.
+ - Negative indices have been wrapped appropriately.
+
+ This function also manipulates ``arr`` if necessary, inserting and/or transposing dimensions as dictated
+ by advanced indexing rules. Finally, it calculates the output shape, new split axis, and balanced status
+ of the resulting indexed array.
+
+ Parameters
+ ----------
+ arr : DNDarray
+ The ``DNDarray`` to be indexed.
+ key : array-like indexer
+ The raw key used for indexing.
+ return_local_indices : bool, optional
+ Whether to map the split-axis indices from global to process-local indices. This is only applied
+ when the indexing key along the split dimension is ordered (i.e., ``split_key_is_ordered == 1``).
+ Default: ``False``.
+ op : str, optional
+ The indexing context for which the key is being processed. Can be ``"get"`` for ``__getitem__``
+ or ``"set"`` for ``__setitem__``. Default: ``None``.
+
+ Returns
+ -------
+ tuple
+ A tuple containing two elements: ``(arr, processed_key)``.
+
+ - arr (DNDarray):
+ The array to be indexed. Its dimensions may have been transposed or expanded if advanced,
+ dimensional, or broadcasted indexing was used.
+ - processed_key (ProcessedKey):
+ A named tuple containing the resolved state required to execute the indexing operation,
+ consisting of the following fields:
+
+ - key (tuple): The processed, Torch-compatible index. Note: Indices along the split axis
+ are local if ordered indexing is used, but remain global if unordered indexing is required.
+ - op_type (str): The categorized indexing routing (``"scalar"``, ``"slice"``,
+ ``"descending_slice"``, ``"distr_mask"``, ``"local_mask"``, ``"advanced"``, or ``"distributed"``).
+ - output_shape (tuple): The global shape of the resulting array.
+ - output_split (int or None): The split axis of the resulting array.
+ - split_key_is_ordered (int): Monotonicity of the split key (``1``: ascending, ``0``: unordered,
+ ``-1``: descending).
+ - key_is_mask_like (bool): Whether the key acts as a boolean mask.
+ - out_is_balanced (bool): Whether the resulting ``DNDarray`` maintains load balance.
+ - root (int or None): The root MPI process ID if single-element indexing along the split
+ axis isolate data to one rank.
+ - backwards_transpose_axes (tuple): The axes required to transpose ``arr`` back to its
+ original shape if advanced indexing triggered a transposition.
+ """
+ # early out for scalar key
+ is_scalar = np.isscalar(key) or getattr(key, "ndim", 1) == 0
+
+ is_boolean = isinstance(key, bool) or (
+ hasattr(key, "dtype")
+ and key.dtype in (ht_bool, ht_uint8, torch.bool, torch.uint8, np.bool_, np.uint8)
+ )
+
+ if is_scalar and not is_boolean:
+ if arr.ndim == 0 and op == "get":
+ raise IndexError(
+ "Too many indices for DNDarray: DNDarray is 0-dimensional, but 1 were indexed"
+ )
+
+ output_shape = arr.gshape[1:]
+ output_split = None if arr.split in (None, 0) else arr.split - 1
+ key, root = _process_scalar_key(
+ arr, key, indexed_axis=0, return_local_indices=return_local_indices
+ )
+
+ return arr, ProcessedKey(
+ key=key,
+ op_type="scalar",
+ output_shape=tuple(output_shape),
+ output_split=output_split,
+ split_key_is_ordered=1,
+ key_is_mask_like=False,
+ out_is_balanced=True,
+ root=root,
+ backwards_transpose_axes=tuple(range(arr.ndim)),
+ )
+
+ # cast any numpy keys to torch tensor
+ if isinstance(key, np.ndarray):
+ key = torch.from_numpy(key)
+
+ # evaluate if this is a distributed fast-path mask before we modify the key
+ distr_mask_fast_path = False
+ # mask along split axis within tuple?
+ if arr.is_distributed():
+ if isinstance(key, tuple) and len(key) > arr.split:
+ split_key = key[arr.split]
+ elif isinstance(key, DNDarray):
+ split_key = key
+ else:
+ split_key = None
+
+ if (
+ isinstance(split_key, DNDarray)
+ and split_key.dtype in (ht_bool, ht_uint8)
+ and split_key.split == arr.split
+ ):
+ # exact shape match
+ if split_key.gshape == arr.gshape:
+ # "get" flattens to 1D
+ # if split > 0, local flattening scrambles global C-order
+ if op == "set" or (op == "get" and arr.split == 0):
+ distr_mask_fast_path = True
+ elif (
+ split_key.ndim == 1
+ and arr.split == 0
+ and split_key.gshape == (arr.gshape[arr.split],)
+ ):
+ # 1D mask on split=0
+ distr_mask_fast_path = True
+
+ # early out if mask and not tuple key
+ if distr_mask_fast_path and not isinstance(key, tuple):
+ return arr, ProcessedKey(
+ key=key.larray,
+ op_type="distr_mask",
+ output_shape=(), # Dummy shape, bypassed safely in __setitem__
+ output_split=0 if op == "get" else arr.split,
+ split_key_is_ordered=0,
+ key_is_mask_like=True,
+ out_is_balanced=False,
+ root=None,
+ backwards_transpose_axes=tuple(range(arr.ndim)),
+ )
+
+ # normalize index components
+ if isinstance(key, DNDarray):
+ if key.dtype not in (ht_bool, ht_uint8) and key.split is None:
+ key = key.larray.to(torch.int64)
+ elif isinstance(key, (list, tuple)):
+ key = type(key)(
+ k.larray.to(torch.int64)
+ if isinstance(k, DNDarray) and k.dtype not in (ht_bool, ht_uint8) and k.split is None
+ else k
+ for k in key
+ )
+
+ # 1D boolean mask resolution
+ first = key[0] if isinstance(key, tuple) and len(key) >= 1 else key
+ if isinstance(first, (DNDarray, torch.Tensor)) and arr.ndim >= 1:
+ first_dtype = getattr(first, "dtype", None)
+ first_ndim = getattr(first, "ndim", 0)
+ first_shape = tuple(getattr(first, "shape", ()))
+
+ if (
+ not distr_mask_fast_path
+ and first_ndim == 1
+ and first_shape == (arr.gshape[0],)
+ and first_dtype in (ht_bool, ht_uint8, torch.bool, torch.uint8)
+ ):
+ if isinstance(first, DNDarray):
+ nz = first.nonzero()
+ if isinstance(nz, tuple):
+ nz = nz[0]
+ if getattr(nz, "ndim", 1) > 1 and nz.shape[-1] == 1:
+ nz = nz.squeeze(-1)
+ idx0 = nz
+ elif isinstance(first, torch.Tensor):
+ idx0 = torch.nonzero(first, as_tuple=False).flatten()
+ else:
+ raise Exception(f"Unexpected type {type(first)}")
+
+ key = (idx0,) + key[1:] if isinstance(key, tuple) else (idx0,)
+
+ output_shape = list(arr.gshape)
+ split_bookkeeping = [None] * arr.ndim
+ new_split = arr.split
+ arr_is_distributed = False
+ if arr.split is not None:
+ split_bookkeeping[arr.split] = "split"
+ if arr.is_distributed():
+ counts, displs = arr.counts_displs()
+ arr_is_distributed = True
+
+ advanced_indexing = False
+ split_key_is_ordered = 1
+ key_is_mask_like = False
+ out_is_balanced = True if not arr.is_distributed() else arr.balanced
+ root = None
+ backwards_transpose_axes = tuple(range(arr.ndim))
+
+ if isinstance(key, list):
+ try:
+ key = torch.tensor(key, device=arr.larray.device)
+ except RuntimeError:
+ raise IndexError("Invalid indices: expected a list of integers, got {}".format(key))
+
+ if isinstance(key, (DNDarray, torch.Tensor)):
+ if key.dtype in (ht_bool, ht_uint8, torch.bool, torch.uint8):
+ # boolean indexing: shape must be consistent with arr.shape
+ key_ndim = key.ndim
+ if not tuple(key.shape) == arr.shape[:key_ndim]:
+ raise IndexError(
+ "Boolean index of shape {} does not match indexed array of shape {}".format(
+ tuple(key.shape), arr.shape
+ )
+ )
+ if key_ndim == 0:
+ # 0-D boolean mask: keep as 0-D tensor, do not extract non-zero
+ key = key.larray if isinstance(key, DNDarray) else key
+ else:
+ # extract non-zero elements
+ try:
+ key = key.nonzero(as_tuple=True)
+ except TypeError:
+ key = key.nonzero()
+
+ key_is_mask_like = True
+ else:
+ # advanced indexing on first dimension: first dim will expand to shape of key
+ output_shape = tuple(list(key.shape) + output_shape[1:])
+ # adjust split axis accordingly
+ if arr_is_distributed:
+ if arr.split != 0:
+ # split axis is not affected
+ split_bookkeeping = [None] * key.ndim + split_bookkeeping[1:]
+ new_split = (
+ split_bookkeeping.index("split") if "split" in split_bookkeeping else None
+ )
+ out_is_balanced = arr.balanced
+ else:
+ # split axis is affected
+ if key.ndim > 1:
+ key_numel = key.numel()
+ if key_numel == arr.shape[0]:
+ new_split = tuple(key.shape).index(arr.shape[0])
+ else:
+ new_split = key.ndim - 1
+ else:
+ new_split = 0
+
+ key_is_dist = isinstance(key, DNDarray) and key.is_distributed()
+ if isinstance(key, DNDarray):
+ out_is_balanced = key.balanced
+ key = key.larray
+ else:
+ out_is_balanced = True
+
+ # normalize negative indices
+ if key.dtype in (torch.int8, torch.int16, torch.int32, torch.int64):
+ dim = arr.gshape[0]
+ if ((key < -dim) | (key >= dim)).any():
+ raise IndexError(f"index out of bounds for axis 0 with size {dim}")
+ key = torch.where(key < 0, key + dim, key)
+
+ # identify ordered key
+ if key_is_dist or key.ndim > 1:
+ split_key_is_ordered = 0
+ else:
+ try:
+ sorted_k, _ = torch.sort(key, stable=True)
+ except TypeError:
+ sorted_k, _ = torch.sort(key)
+ split_key_is_ordered = int((key == sorted_k).all().item())
+
+ # unordered local keys
+ if not split_key_is_ordered and not key_is_dist:
+ if op == "get":
+ # prepare for distributed non-ordered indexing: distribute local key
+ key = factories.array(key, split=new_split, device=arr.device).larray
+ out_is_balanced = True
+ else:
+ out_is_balanced = True
+
+ # ordered keys
+ if split_key_is_ordered:
+ # extract local key
+ cond1 = key >= displs[arr.comm.rank]
+ cond2 = key < displs[arr.comm.rank] + counts[arr.comm.rank]
+ key = key[cond1 & cond2]
+ if return_local_indices:
+ key -= displs[arr.comm.rank]
+ out_is_balanced = False
+ else:
+ try:
+ out_is_balanced = key.balanced
+ new_split = key.split
+ key = key.larray
+ except AttributeError:
+ # torch key, non-distributed indexed array
+ out_is_balanced = True
+ new_split = None
+
+ # define indexing type
+ if root is not None:
+ op_type = "scalar"
+ elif split_key_is_ordered == 0:
+ op_type = "distributed"
+ elif key_is_mask_like:
+ op_type = "local_mask"
+ else:
+ op_type = "advanced"
+
+ return arr, ProcessedKey(
+ key=key,
+ op_type=op_type,
+ output_shape=tuple(output_shape),
+ output_split=new_split,
+ split_key_is_ordered=split_key_is_ordered,
+ key_is_mask_like=key_is_mask_like,
+ out_is_balanced=out_is_balanced,
+ root=root,
+ backwards_transpose_axes=backwards_transpose_axes,
+ )
+
+ if isinstance(key, (tuple, list)):
+ key = list(key)
+ else:
+ key = [key]
+
+ # check for ellipsis, newaxis. NB: (np.newaxis is None)==True
+ def is_0d_bool(k):
+ if isinstance(k, bool):
+ return True
+ if hasattr(k, "dtype") and k.dtype in (
+ ht_bool,
+ ht_uint8,
+ torch.bool,
+ torch.uint8,
+ ):
+ if getattr(k, "ndim", 1) == 0:
+ return True
+ return False
+
+ add_dims = sum(k is None or is_0d_bool(k) for k in key)
+ ellipsis = sum(isinstance(k, type(...)) for k in key)
+ if ellipsis > 1:
+ raise ValueError("indexing key can only contain 1 Ellipsis (...)")
+ if ellipsis:
+ # key contains exactly 1 ellipsis
+ # replace with explicit `slice(None)` for affected dimensions
+ # output_shape, split_bookkeeping not affected
+ expand_key = [slice(None)] * (arr.ndim + add_dims)
+ ellipsis_index = key.index(...)
+ ellipsis_dims = arr.ndim - (len(key) - ellipsis - add_dims)
+ expand_key[:ellipsis_index] = key[:ellipsis_index]
+ expand_key[ellipsis_index + ellipsis_dims :] = key[ellipsis_index + 1 :]
+ key = expand_key
+ while add_dims > 0:
+ # expand array dims: output_shape, split_bookkeeping to reflect newaxis
+ # replace newaxis with slice(None), replace 0-D bools with a target slice
+ for i, k in reversed(list(enumerate(key))):
+ if k is None or is_0d_bool(k):
+ if k is None:
+ key[i] = slice(None)
+ else:
+ val = bool(k.item() if hasattr(k, "item") else k)
+ key[i] = slice(None) if val else slice(0, 0)
+
+ arr = arr.expand_dims(i - add_dims + 1)
+ output_shape = (
+ output_shape[: i - add_dims + 1] + [1] + output_shape[i - add_dims + 1 :]
+ )
+ split_bookkeeping = (
+ split_bookkeeping[: i - add_dims + 1]
+ + [None]
+ + split_bookkeeping[i - add_dims + 1 :]
+ )
+ add_dims -= 1
+
+ # recalculate new_split, transpose_axes after dimensions manipulation
+ new_split = split_bookkeeping.index("split") if "split" in split_bookkeeping else None
+ transpose_axes, backwards_transpose_axes = tuple(range(arr.ndim)), tuple(range(arr.ndim))
+ # check for advanced indexing and slices
+ advanced_indexing_dims = []
+ advanced_indexing_shapes = []
+ lose_dims = 0
+ for i, k in enumerate(key):
+ if isinstance(k, DNDarray) and k.ndim == 0:
+ k = k.larray.item()
+ key[i] = k
+ # for robustness: handle list/tuple keys that contain DNDarrays
+ elif isinstance(k, (list, tuple)) and any(isinstance(kk, DNDarray) for kk in k):
+ # Case 1: singleton container (common from where/nonzero): (idx,) -> idx
+ if len(k) == 1 and isinstance(k[0], DNDarray):
+ k = k[0]
+ key[i] = k
+
+ else:
+ # Case 2: sequence of scalar DNDarrays -> unwrap to python scalars
+ new_k = []
+ all_scalar = True
+ for kk in k:
+ if isinstance(kk, DNDarray):
+ if kk.ndim != 0:
+ all_scalar = False
+ break
+ new_k.append(kk.larray.item())
+ else:
+ new_k.append(kk)
+
+ if all_scalar:
+ k = new_k
+ key[i] = k
+ else:
+ # This is an ambiguous nested "tuple of index arrays" inside a single axis.
+ # In NumPy semantics such tuples belong at TOP LEVEL (arr[idx0, idx1, ...]),
+ # not nested as one axis key.
+ raise TypeError(
+ "Nested tuple/list of non-scalar DNDarray indices is not supported. "
+ "Pass them as separate indices (e.g. arr[idx0, idx1, ...]) or unwrap "
+ "singleton tuples (e.g. idx = idx[0])."
+ )
+
+ if np.isscalar(k) or getattr(k, "ndim", 1) == 0:
+ # single-element indexing along axis i
+ try:
+ output_shape[i], split_bookkeeping[i] = None, None
+ except IndexError:
+ raise IndexError(
+ f"Too many indices for DNDarray: DNDarray is {arr.ndim}-dimensional, but {len(key)} dimensions were indexed"
+ )
+ lose_dims += 1
+ if i == arr.split:
+ key[i], root = _process_scalar_key(
+ arr, k, indexed_axis=i, return_local_indices=return_local_indices
+ )
+ else:
+ key[i], _ = _process_scalar_key(arr, k, indexed_axis=i, return_local_indices=False)
+ elif isinstance(k, Iterable) or isinstance(k, DNDarray):
+ advanced_indexing = True
+ advanced_indexing_dims.append(i)
+
+ is_fast_path_component = distr_mask_fast_path and i == arr.split
+
+ if is_fast_path_component:
+ key[i] = k.larray if isinstance(k, DNDarray) else k
+ advanced_indexing_shapes.append(tuple(k.shape))
+ # skip the rest, local boolean masking along split axis
+ continue
+
+ if not isinstance(k, DNDarray):
+ k = factories.array(k, device=arr.device, comm=arr.comm, copy=None)
+
+ # normalize negative integer indices (NumPy/PyTorch semantics) and validate bounds
+ if k.dtype in (types.int32, types.int64) and k.ndim >= 1:
+ dim = arr.gshape[i]
+
+ # compute local flags even if k.larray is empty (any() on empty -> False)
+ invalid_local = ((k.larray < -dim) | (k.larray >= dim)).any().item()
+ has_neg_local = (k.larray < 0).any().item()
+
+ # Decide once, then ALL ranks take the same path for collectives
+ do_reduce = (
+ arr.comm is not None and getattr(arr.comm, "size", 1) > 1 and k.is_distributed()
+ )
+
+ if do_reduce:
+ invalid_sum = arr.comm.allreduce(int(invalid_local), op=MPI.SUM)
+ has_neg_sum = arr.comm.allreduce(int(has_neg_local), op=MPI.SUM)
+ else:
+ invalid_sum = int(invalid_local)
+ has_neg_sum = int(has_neg_local)
+
+ if invalid_sum > 0:
+ raise IndexError(f"index out of bounds for axis {i} with size {dim}")
+
+ if has_neg_sum > 0:
+ k_l = k.larray.clone()
+ k_l[k_l < 0] += dim
+ k = factories.array(
+ k_l,
+ dtype=k.dtype,
+ split=k.split,
+ device=arr.device,
+ comm=arr.comm,
+ copy=False,
+ )
+
+ advanced_indexing_shapes.append(k.gshape)
+ if arr_is_distributed and i == arr.split:
+ if (
+ not k.is_distributed()
+ and k.ndim == 1
+ and (k.larray == torch.sort(k.larray, stable=True)[0]).all()
+ ):
+ split_key_is_ordered = 1
+ out_is_balanced = False
+ else:
+ split_key_is_ordered = 0
+
+ # redistribute key along last axis to match split axis of indexed array
+ k = k.resplit(-1)
+ out_is_balanced = True
+ key[i] = k
+
+ elif isinstance(k, slice) and k != slice(None):
+ if k.step == 0:
+ raise ValueError("Slice step cannot be zero")
+ start, stop, step = slice(k.start, k.stop, k.step).indices(arr.gshape[i])
+
+ if step < 0 and start > stop:
+ # PyTorch doesn't support negative step
+ key[i] = torch.arange(
+ start, stop, step, device=arr.larray.device, dtype=torch.int64
+ )
+ output_shape[i] = len(key[i])
+
+ if arr_is_distributed and new_split == i:
+ split_key_is_ordered = -1
+ # flip key and keep process-local indices
+ key[i] = key[i].flip(0)
+ cond1 = key[i] >= displs[arr.comm.rank]
+ cond2 = key[i] < displs[arr.comm.rank] + counts[arr.comm.rank]
+ key[i] = key[i][cond1 & cond2]
+ if return_local_indices:
+ key[i] -= displs[arr.comm.rank]
+ # slices can result in unbalanced chunks
+ out_is_balanced = False
+
+ elif step > 0 and start < stop:
+ output_shape[i] = len(range(start, stop, step))
+
+ if arr_is_distributed and new_split == i:
+ split_key_is_ordered = 1
+ out_is_balanced = False
+ local_arr_end = displs[arr.comm.rank] + counts[arr.comm.rank]
+ if stop > displs[arr.comm.rank] and start < local_arr_end:
+ index_in_cycle = (displs[arr.comm.rank] - start) % step
+ if start >= displs[arr.comm.rank]:
+ # slice begins on current rank
+ local_start = start - displs[arr.comm.rank]
+ else:
+ local_start = 0 if index_in_cycle == 0 else step - index_in_cycle
+ if stop <= local_arr_end:
+ # slice ends on current rank
+ local_stop = stop - displs[arr.comm.rank]
+ else:
+ local_stop = counts[arr.comm.rank]
+
+ key[i] = slice(local_start, local_stop, step)
+ else:
+ key[i] = slice(0, 0)
+ elif step == 0:
+ raise ValueError("Slice step cannot be zero")
+ else:
+ key[i] = slice(0, 0)
+ output_shape[i] = 0
+
+ if advanced_indexing:
+ # adv indexing key elements are DNDarrays: extract torch tensors
+ # options: 1. key is mask-like (covers boolean mask as well), 2. adv indexing along split axis, 3. everything else
+ # 1. define key as mask-like if each element of key is a DNDarray, and all elements of key are of the same shape, and the advanced-indexing dimensions are consecutive
+ key_is_mask_like = key_is_mask_like or (
+ len(advanced_indexing_dims) > 1
+ and all(isinstance(k, DNDarray) for k in key)
+ and len(set(k.shape for k in key)) == 1
+ and torch.tensor(advanced_indexing_dims).diff().eq(1).all().item()
+ )
+ # if split axis is affected by advanced indexing, keep track of non-split dimensions for later
+ if arr.is_distributed() and arr.split in advanced_indexing_dims:
+ non_split_dims = list(advanced_indexing_dims).copy()
+ if arr.split is not None:
+ non_split_dims.remove(arr.split)
+ # 1. key is mask-like
+ if key_is_mask_like:
+ key = list(key)
+ key_splits = [k.split for k in key]
+ if arr.split is not None and arr.split in advanced_indexing_dims:
+ split_key_pos = advanced_indexing_dims.index(arr.split)
+
+ if not key_splits.count(key_splits[split_key_pos]) == len(key_splits):
+ if (
+ key_splits[arr.split] is not None
+ and key_splits.count(None) == len(key_splits) - 1
+ ):
+ for i in non_split_dims:
+ key[i] = factories.array(
+ key[i],
+ split=key_splits[arr.split],
+ device=arr.device,
+ comm=arr.comm,
+ copy=None,
+ )
+ else:
+ raise IndexError(
+ f"Indexing arrays must be distributed along the same dimension, got splits {key_splits}."
+ )
+ else:
+ # all key_splits must be the same, otherwise raise IndexError
+ if not key_splits.count(key_splits[0]) == len(key_splits):
+ raise IndexError(
+ f"Indexing arrays must be distributed along the same dimension, got splits {key_splits}."
+ )
+ # all key elements are now DNDarrays of the same shape, same split axis
+ # 2. advanced indexing along split axis
+ if arr.is_distributed() and arr.split in advanced_indexing_dims:
+ if distr_mask_fast_path:
+ # mask is already a local tensor, just extract any other advanced indices
+ for i in non_split_dims:
+ if isinstance(key[i], DNDarray):
+ key[i] = key[i].larray
+ elif split_key_is_ordered == 1:
+ # extract torch tensors, keep process-local indices only
+ k = key[arr.split].larray
+ cond1 = k >= displs[arr.comm.rank]
+ cond2 = k < displs[arr.comm.rank] + counts[arr.comm.rank]
+ k = k[cond1 & cond2]
+ if return_local_indices:
+ k -= displs[arr.comm.rank]
+ key[arr.split] = k
+ for i in non_split_dims:
+ if key_is_mask_like:
+ # select the same elements along non-split dimensions
+ key[i] = key[i].larray[cond1 & cond2]
+ else:
+ key[i] = key[i].larray
+ elif split_key_is_ordered == 0:
+ # extract torch tensors, any other communication + mask-like case are handled in __getitem__ or __setitem__
+ for i in advanced_indexing_dims:
+ key[i] = key[i].larray
+ # split_key_is_ordered == -1 not treated here as it is slicing, not advanced indexing
+ else:
+ # advanced indexing does not affect split axis, return torch tensors
+ for i in advanced_indexing_dims:
+ key[i] = key[i].larray
+ # all adv indexing keys are now torch tensors
+
+ # shapes of adv indexing arrays must be broadcastable
+ try:
+ broadcasted_shape = torch.broadcast_shapes(*advanced_indexing_shapes)
+ except RuntimeError:
+ raise IndexError(
+ "Shape mismatch: indexing arrays could not be broadcast together with shapes: {}".format(
+ advanced_indexing_shapes
+ )
+ )
+ add_dims = len(broadcasted_shape) - len(advanced_indexing_dims)
+ if (
+ len(advanced_indexing_dims) == 1
+ or list(range(advanced_indexing_dims[0], advanced_indexing_dims[-1] + 1))
+ == advanced_indexing_dims
+ ):
+ # dimensions affected by advanced indexing are consecutive:
+ output_shape[
+ advanced_indexing_dims[0] : advanced_indexing_dims[0] + len(advanced_indexing_dims)
+ ] = broadcasted_shape
+ if key_is_mask_like:
+ # advanced indexing dimensions will be collapsed into one dimension
+ if (
+ "split" in split_bookkeeping
+ and split_bookkeeping.index("split") in advanced_indexing_dims
+ ):
+ split_bookkeeping[
+ advanced_indexing_dims[0] : advanced_indexing_dims[0]
+ + len(advanced_indexing_dims)
+ ] = ["split"]
+ else:
+ split_bookkeeping[
+ advanced_indexing_dims[0] : advanced_indexing_dims[0]
+ + len(advanced_indexing_dims)
+ ] = [None]
+ else:
+ # Replace the original slice with a properly sized list representing the broadcasted shape
+ # adv_sb = slice of split_bookkeeping corresponding to advanced_indexing_dims
+ adv_sb = split_bookkeeping[
+ advanced_indexing_dims[0] : advanced_indexing_dims[-1] + 1
+ ]
+ new_adv_sb = [None] * len(broadcasted_shape)
+ if "split" in adv_sb:
+ # track 'split', adjust for added dimensions
+ new_idx = adv_sb.index("split") + add_dims
+ if new_idx < 0:
+ new_idx = 0
+ new_adv_sb[new_idx] = "split"
+
+ split_bookkeeping = (
+ split_bookkeeping[: advanced_indexing_dims[0]]
+ + new_adv_sb
+ + split_bookkeeping[advanced_indexing_dims[-1] + 1 :]
+ )
+ else:
+ # advanced-indexing dimensions are not consecutive:
+ # transpose array to make the advanced-indexing dimensions consecutive as the first dimensions
+ non_adv_ind_dims = list(i for i in range(arr.ndim) if i not in advanced_indexing_dims)
+ # keep track of transpose axes order, to be able to transpose back later
+ transpose_axes = tuple(advanced_indexing_dims + non_adv_ind_dims)
+ arr = arr.transpose(transpose_axes)
+ backwards_transpose_axes = tuple(
+ torch.tensor(transpose_axes, device=arr.larray.device).argsort(stable=True).tolist()
+ )
+ # output shape and split bookkeeping
+ output_shape = list(output_shape[i] for i in transpose_axes)
+ output_shape[: len(advanced_indexing_dims)] = broadcasted_shape
+
+ split_bookkeeping = list(split_bookkeeping[i] for i in transpose_axes)
+ adv_sb = split_bookkeeping[: len(advanced_indexing_dims)]
+ new_adv_sb = [None] * len(broadcasted_shape)
+
+ if "split" in adv_sb:
+ new_idx = adv_sb.index("split") + add_dims
+ if new_idx < 0:
+ new_idx = 0
+ new_adv_sb[new_idx] = "split"
+
+ split_bookkeeping = new_adv_sb + split_bookkeeping[len(advanced_indexing_dims) :]
+
+ # modify key to match the new dimension order
+ key = [key[i] for i in advanced_indexing_dims] + [key[i] for i in non_adv_ind_dims]
+ # update advanced-indexing dims
+ advanced_indexing_dims = list(range(len(advanced_indexing_dims)))
+
+ # expand key to match the number of dimensions of the DNDarray
+ if arr.ndim > len(key):
+ key += [slice(None)] * (arr.ndim - len(key))
+
+ key = tuple(key)
+ for i in range(output_shape.count(None)):
+ lost_dim = output_shape.index(None)
+ output_shape.remove(None)
+ split_bookkeeping = split_bookkeeping[:lost_dim] + split_bookkeeping[lost_dim + 1 :]
+ output_shape = tuple(output_shape)
+ new_split = split_bookkeeping.index("split") if "split" in split_bookkeeping else None
+
+ if root is not None:
+ op_type = "scalar"
+ elif split_key_is_ordered == 0:
+ op_type = "distributed"
+ elif split_key_is_ordered == -1:
+ op_type = "descending_slice"
+ elif key_is_mask_like:
+ op_type = "distr_mask" if distr_mask_fast_path else "local_mask"
+ else:
+ op_type = "advanced"
+
+ return arr, ProcessedKey(
+ key=tuple(key),
+ op_type=op_type,
+ output_shape=tuple(output_shape),
+ output_split=new_split,
+ split_key_is_ordered=split_key_is_ordered,
+ key_is_mask_like=key_is_mask_like,
+ out_is_balanced=out_is_balanced,
+ root=root,
+ backwards_transpose_axes=backwards_transpose_axes,
+ )
+
+
class DNDarray:
"""
- Distributed N-Dimensional array. The core element of HeAT. It is composed of
+ Distributed N-Dimensional array. The core element of Heat. It is composed of
PyTorch tensors local to each process.
Parameters
----------
array : torch.Tensor
Local array elements
- gshape : Tuple[int,...]
+ gshape : tuple[int,...]
The global shape of the array
dtype : datatype
The datatype of the array
@@ -64,9 +976,9 @@ class DNDarray:
def __init__(
self,
array: torch.Tensor,
- gshape: Tuple[int, ...],
+ gshape: tuple[int, ...],
dtype: datatype,
- split: Union[int, None],
+ split: int | None,
device: Device,
comm: Communication,
balanced: bool,
@@ -77,10 +989,10 @@ def __init__(
self.__split = split
self.__device = device
self.__comm = comm
- self.__balanced = balanced
+ self.__balanced: bool = balanced
self.__ishalo = False
- self.__halo_next = None
- self.__halo_prev = None
+ self.__halo_next: torch.Tensor | None = None
+ self.__halo_prev: torch.Tensor | None = None
self.__partitions_dict__ = None
self.__lshape_map = None
@@ -116,7 +1028,7 @@ def dtype(self) -> datatype:
return self.__dtype
@property
- def gshape(self) -> Tuple:
+ def gshape(self) -> tuple:
"""
Returns the global shape of the ``DNDarray`` across all processes
"""
@@ -263,7 +1175,7 @@ def lnumel(self) -> int:
return np.prod(self.__array.shape)
@property
- def lloc(self) -> Union[DNDarray, None]:
+ def lloc(self) -> "DNDarray" | None:
"""
Local item setter and getter. i.e. this function operates on a local
level and only on the PyTorch tensors composing the :class:`DNDarray`.
@@ -272,7 +1184,7 @@ def lloc(self) -> Union[DNDarray, None]:
Parameters
----------
- key : int or slice or Tuple[int,...]
+ key : int or slice or tuple[int,...]
Indices of the desired data.
value : scalar, optional
All types compatible with pytorch tensors, if none given then this is a getter function
@@ -297,7 +1209,7 @@ def lloc(self) -> Union[DNDarray, None]:
return LocalIndex(self.__array)
@property
- def lshape(self) -> Tuple[int]:
+ def lshape(self) -> tuple[int]:
"""
Returns the shape of the ``DNDarray`` on each node
"""
@@ -318,36 +1230,35 @@ def real(self) -> DNDarray:
return complex_math.real(self)
@property
- def shape(self) -> Tuple[int]:
+ def shape(self) -> tuple[int, ...]:
"""
Returns the shape of the ``DNDarray`` as a whole
"""
return self.__gshape
@property
- def split(self) -> int:
+ def split(self) -> int | None:
"""
Returns the axis on which the ``DNDarray`` is split
"""
return self.__split
- @property
- def stride(self) -> Tuple[int]:
+ def stride(self) -> tuple[int, ...]:
"""
Returns the steps in each dimension when traversing a ``DNDarray``. torch-like usage: ``self.stride()``
"""
- return self.__array.stride
+ return self.__array.stride()
@property
- def strides(self) -> Tuple[int]:
+ def strides(self) -> tuple[int, ...]:
"""
Returns bytes to step in each dimension when traversing a ``DNDarray``. numpy-like usage: ``self.strides()``
"""
- steps = list(self.larray.stride())
+ steps = list(self.__array.stride())
try:
- itemsize = self.larray.untyped_storage().element_size()
+ itemsize = self.__array.untyped_storage().element_size()
except AttributeError:
- itemsize = self.larray.storage().element_size()
+ itemsize = self.__array.storage().element_size()
strides = tuple(step * itemsize for step in steps)
return strides
@@ -543,7 +1454,13 @@ def astype(self, dtype, copy=True) -> DNDarray:
casted_array = self.__array.type(dtype.torch_type())
if copy:
return DNDarray(
- casted_array, self.shape, dtype, self.split, self.device, self.comm, self.balanced
+ casted_array,
+ gshape=self.shape,
+ dtype=dtype,
+ split=self.split,
+ device=self.device,
+ comm=self.comm,
+ balanced=self.balanced,
)
self.__array = casted_array
@@ -551,7 +1468,7 @@ def astype(self, dtype, copy=True) -> DNDarray:
return self
- def balance_(self) -> DNDarray:
+ def balance_(self) -> None:
"""
Function for balancing a :class:`DNDarray` between all nodes. To determine if this is needed use the :func:`is_balanced()` function.
If the ``DNDarray`` is already balanced this function will do nothing. This function modifies the ``DNDarray``
@@ -587,6 +1504,8 @@ def balance_(self) -> DNDarray:
[1/2] (7, 2) (2, 2)
[2/2] (7, 2) (2, 2)
"""
+ if not self.is_distributed():
+ self.__balanced = True
if self.is_balanced(force_check=True):
return
self.redistribute_()
@@ -597,7 +1516,7 @@ def __bool__(self) -> bool:
"""
return self.__cast(bool)
- def __cast(self, cast_function) -> Union[float, int]:
+ def __cast(self, cast_function) -> float | int:
"""
Implements a generic cast function for ``DNDarray`` objects.
@@ -623,7 +1542,7 @@ def __cast(self, cast_function) -> Union[float, int]:
raise TypeError("only size-1 arrays can be converted to Python scalars")
- def collect_(self, target_rank: Optional[int] = 0) -> None:
+ def collect_(self, target_rank: int | None = 0) -> None:
"""
A method collecting a distributed DNDarray to one MPI rank, chosen by the `target_rank` variable.
It is a specific case of the ``redistribute_`` method.
@@ -676,7 +1595,7 @@ def __complex__(self) -> DNDarray:
"""
return self.__cast(complex)
- def counts_displs(self) -> Tuple[Tuple[int], Tuple[int]]:
+ def counts_displs(self) -> tuple[tuple[int], tuple[int]]:
"""
Returns actual counts (number of items per process) and displacements (offsets) of the DNDarray.
Does not assume load balance.
@@ -685,8 +1604,8 @@ def counts_displs(self) -> Tuple[Tuple[int], Tuple[int]]:
counts = self.lshape_map[:, self.split]
displs = [0] + torch.cumsum(counts, dim=0)[:-1].tolist()
return tuple(counts.tolist()), tuple(displs)
- else:
- raise ValueError("Non-distributed DNDarray. Cannot calculate counts and displacements.")
+
+ raise ValueError("Non-distributed DNDarray. Cannot calculate counts and displacements.")
def cpu(self) -> DNDarray:
"""
@@ -714,9 +1633,10 @@ def create_lshape_map(self, force_check: bool = False) -> torch.Tensor:
lshape_map = torch.zeros(
(self.comm.size, self.ndim), dtype=torch.int64, device=self.device.torch_device
)
- if not self.is_distributed:
+ if not self.is_distributed():
lshape_map[:] = torch.tensor(self.gshape, device=self.device.torch_device)
- return lshape_map
+ self.__lshape_map = lshape_map
+ return lshape_map.clone()
if self.is_balanced(force_check=True):
for i in range(self.comm.size):
_, lshape, _ = self.comm.chunk(self.gshape, self.split, rank=i)
@@ -789,15 +1709,15 @@ def create_partition_interface(self):
part_tiling = [1] * self.ndim
lcls = [0] * self.ndim
- z = torch.tensor([0], device=self.device.torch_device, dtype=self.dtype.torch_type())
+ z = torch.tensor([0], device=self.device.torch_device, dtype=torch.int64)
+
if self.split is not None:
starts = torch.cat((z, torch.cumsum(lshape_map[:, self.split], dim=0)[:-1]), dim=0)
lcls[self.split] = self.comm.rank
part_tiling[self.split] = self.comm.size
+ start_idx_map[:, self.split] = starts
else:
- starts = torch.zeros(self.ndim, dtype=torch.int, device=self.device.torch_device)
-
- start_idx_map[:, self.split] = starts
+ start_idx_map[:] = 0
partitions = {}
base_key = [0] * self.ndim
@@ -878,18 +1798,499 @@ def fill_diagonal(self, value: float) -> DNDarray:
return self
- def __getitem__(self, key: Union[int, Tuple[int, ...], List[int, ...]]) -> DNDarray:
+ def __broadcast_value(
+ self,
+ key: int | tuple[int, ...] | slice,
+ value: "DNDarray",
+ **kwargs,
+ ):
+ """
+ Broadcasts the assignment DNDarray `value` to the shape of the indexed array `arr[key]` if necessary.
+ """
+ is_scalar = (
+ np.isscalar(value)
+ or getattr(value, "ndim", 1) == 0
+ or (value.shape == (1,) and value.split is None)
+ )
+ if is_scalar:
+ # no need to broadcast
+ return value, is_scalar
+ # need information on indexed array
+ output_shape = kwargs.get("output_shape", None)
+ if output_shape is not None:
+ indexed_dims = len(output_shape)
+ else:
+ if isinstance(key, (int, tuple)):
+ # direct indexing, output_shape has not been calculated
+ # use proxy to avoid MPI communication and limit memory usage
+ indexed_proxy = self.__torch_proxy__()[key]
+ indexed_dims = indexed_proxy.ndim
+ output_shape = tuple(indexed_proxy.shape)
+ else:
+ raise RuntimeError(
+ "Not enough information to broadcast value to indexed array, please provide `output_shape`"
+ )
+ value_shape = value.shape
+ # check if value needs to be broadcasted
+ if value_shape != output_shape:
+ # assess whether the shapes are compatible, starting from the trailing dimension
+ for i in range(1, min(len(value_shape), len(output_shape)) + 1):
+ if i == 1:
+ if value_shape[-i] != output_shape[-i] and not value_shape[-i] == 1:
+ # shapes are not compatible, raise error
+ raise ValueError(
+ f"could not broadcast input array from shape {value_shape} into shape {output_shape}"
+ )
+ else:
+ if value_shape[-i] != output_shape[-i] and (not value_shape[-i] == 1):
+ # shapes are not compatible, raise error
+ raise ValueError(
+ f"could not broadcast input from shape {value_shape} into shape {output_shape}"
+ )
+ # value has more dimensions than indexed array
+ if value.ndim > indexed_dims:
+ # check if all dimensions except the indexed ones are singletons
+ all_singletons = value.shape[: value.ndim - indexed_dims] == (1,) * (
+ value.ndim - indexed_dims
+ )
+ if not all_singletons:
+ raise ValueError(
+ f"could not broadcast input array from shape {value_shape} into shape {output_shape}"
+ )
+ # squeeze out singleton dimensions
+ value = value.squeeze(tuple(range(value.ndim - indexed_dims)))
+ else:
+ while value.ndim < indexed_dims:
+ # broadcasting
+ # expand missing dimensions to align split axis
+ value = value.expand_dims(0)
+ try:
+ value_shape = tuple(torch.broadcast_shapes(value.shape, output_shape))
+ except RuntimeError:
+ raise ValueError(
+ f"could not broadcast input array from shape {value_shape} into shape {output_shape}"
+ )
+ return value, is_scalar
+
+ def __set(
+ self,
+ key: int | tuple[int, ...] | list[int],
+ value: float | "DNDarray" | torch.Tensor,
+ ):
+ """
+ Setter for not advanced indexing, i.e. when arr[key] is an in-place view of arr.
+ """
+ # only assign values if key does not contain empty slices
+ process_is_inactive = self.larray[key].numel() == 0
+ if not process_is_inactive:
+ rhs = value.larray.type(self.dtype.torch_type())
+ key_to_use = key
+
+ # CUDA: make advanced indexing assignment deterministic for duplicate indices
+ if self.larray.is_cuda:
+ key_to_use, rhs = _resolve_duplicate_indices(key_to_use, rhs, self.larray.shape)
+
+ self.larray[key_to_use] = rhs
+ return
+
+ @staticmethod
+ def __advanced_setitem_unordered_local(
+ x_local: torch.Tensor,
+ split_key: torch.Tensor,
+ value_torch: torch.Tensor,
+ *,
+ split_axis: int,
+ value_key_start_dim: int,
+ local_offset: int,
+ local_size: int,
+ value_is_scalar: bool,
+ out_dtype: torch.dtype,
+ base_index: tuple | None = None,
+ ) -> None:
+ """
+ The function is a helper that updates ``x_local`` in-place according to the logical advanced
+ indexing pattern encoded by ``split_key`` and the broadcasted ``value_torch``.
+ This helper operates exclusively on local ``torch.Tensor`` views:
+ - ``x_local`` is the local slice of the distributed array on this rank.
+ - ``split_key`` contains GLOBAL indices along the split axis.
+ - Only those indices that fall into ``[local_offset, local_offset + local_size)``
+ are applied on this rank.
+ """
+ # 1) Local mask: which global indices in `split_key` belong to this rank?
+ global_indices = split_key
+ local_mask = (global_indices >= local_offset) & (global_indices < local_offset + local_size)
+
+ coord = local_mask.nonzero(as_tuple=True)
+
+ if coord[0].numel() == 0:
+ # Nothing to do on this rank, exit early.
+ return
+
+ # 2) Map global → local indices along the split axis
+ global_split_indices = global_indices[coord]
+ local_split_indices = global_split_indices - local_offset
+
+ # build LHS index for x_local (corresponds to self.larray)
+ if base_index is None:
+ lhs_index = [slice(None)] * x_local.ndim
+ else:
+ lhs_index = list(base_index)
+
+ lhs_index[split_axis] = local_split_indices
+ lhs_index = tuple(lhs_index)
+
+ # build RHS index for value_torch
+ if value_is_scalar:
+ rhs = value_torch.to(out_dtype)
+ else:
+ rhs_index = [slice(None)] * value_torch.ndim
+ m = split_key.ndim
+
+ for d in range(m):
+ rhs_index[value_key_start_dim + d] = coord[d]
+
+ rhs = value_torch[tuple(rhs_index)].to(out_dtype)
+
+ if x_local.is_cuda:
+ lhs_index, rhs = _resolve_duplicate_indices(lhs_index, rhs, x_local.shape)
+
+ x_local[lhs_index] = rhs
+
+ def __getitem_scalar(self, p: ProcessedKey) -> DNDarray:
+ """
+ Handles single-element extraction. If the scalar index falls on the
+ split axis, the extracted value is broadcasted from the
+ root process to all others.
+ """
+ if p.root is not None:
+ # Single-element indexing along split axis
+ if self.comm.rank == p.root:
+ indexed_arr = self.larray[p.key]
+ else:
+ indexed_arr = torch.zeros(
+ p.output_shape, dtype=self.larray.dtype, device=self.larray.device
+ )
+ self.comm.Bcast(indexed_arr, root=p.root)
+ else:
+ indexed_arr = self.larray[p.key]
+
+ if self.ndim > 0:
+ self = self.transpose(p.backwards_transpose_axes)
+
+ return DNDarray(
+ indexed_arr,
+ gshape=p.output_shape,
+ dtype=self.dtype,
+ split=p.output_split,
+ device=self.device,
+ comm=self.comm,
+ balanced=p.out_is_balanced,
+ )
+
+ def __getitem_slice(self, p: ProcessedKey) -> "DNDarray":
+ """
+ Handles standard slicing using process-local views. Requires no cross-process
+ MPI communication.
+ """
+ indexed_arr = self.larray[p.key]
+ if self.ndim > 0:
+ self = self.transpose(p.backwards_transpose_axes)
+
+ return DNDarray(
+ indexed_arr,
+ gshape=p.output_shape,
+ dtype=self.dtype,
+ split=p.output_split,
+ device=self.device,
+ comm=self.comm,
+ balanced=p.out_is_balanced,
+ )
+
+ def __getitem_descending_slice_distributed(self, p: ProcessedKey) -> DNDarray:
+ """
+ Handles negative step slicing along the split axis. This is a workaround as torch does not support negative-step slicing.
+ """
+ from .manipulations import flip
+
+ # local indexing
+ indexed_arr = self.larray[p.key]
+ if self.ndim > 0:
+ self = self.transpose(p.backwards_transpose_axes)
+
+ # wrap the reversed local chunks into an unbalanced DNDarray
+ intermediate = DNDarray(
+ indexed_arr,
+ gshape=p.output_shape,
+ dtype=self.dtype,
+ split=p.output_split,
+ device=self.device,
+ comm=self.comm,
+ balanced=False,
+ )
+
+ # global flip to reflect the descending slice
+ return flip(intermediate, axis=p.output_split)
+
+ def __getitem_mask(self, p: ProcessedKey, original_key) -> "DNDarray":
+ """
+ Handles fast-path boolean masking. Applies the mask locally without
+ requiring MPI communication during extraction, returning a flattened array
+ distributed along the specified split axis.
+ """
+ # local masking, then wrap into DNDarray
+ local_mask = p.key
+ local_result = self.larray[local_mask]
+
+ return factories.array(
+ local_result, is_split=p.output_split, device=self.device, comm=self.comm, copy=False
+ )
+
+ def __getitem_advanced_local(self, p: ProcessedKey, original_key) -> "DNDarray":
+ """
+ Handles advanced indexing where no MPI communication is needed
+ (e.g., the split axis is unaffected, or indices are strictly local).
+ """
+ indexed_arr = self.larray[p.key]
+ if self.ndim > 0:
+ self = self.transpose(p.backwards_transpose_axes)
+
+ return DNDarray(
+ indexed_arr,
+ gshape=p.output_shape,
+ dtype=self.dtype,
+ split=p.output_split,
+ device=self.device,
+ comm=self.comm,
+ balanced=p.out_is_balanced,
+ )
+
+ def __getitem_advanced_distributed(self, p: ProcessedKey) -> "DNDarray":
+ """
+ Handles advanced indexing with unordered global indices. Defers to
+ ``__getitem_unordered`` to resolve data dependencies via an ``Alltoallv`` exchange.
+ """
+ self, indexed_arr = self.__getitem_unordered(
+ key=p.key,
+ output_shape=p.output_shape,
+ output_split=p.output_split,
+ out_is_balanced=p.out_is_balanced,
+ key_is_mask_like=p.key_is_mask_like,
+ backwards_transpose_axes=p.backwards_transpose_axes,
+ )
+ return indexed_arr
+
+ def __getitem_unordered(
+ self,
+ key: tuple,
+ output_shape: tuple,
+ output_split: int,
+ out_is_balanced: bool,
+ key_is_mask_like: bool,
+ backwards_transpose_axes: tuple,
+ ) -> DNDarray:
+ """
+ Handles the MPI communication (Alltoallv) when the key along the
+ split axis is unordered and indices are global.
+ """
+ _, displs = self.counts_displs()
+ rank = self.comm.rank
+
+ key_is_single_tensor = isinstance(key, torch.Tensor)
+ split_key = key if key_is_single_tensor else key[self.split]
+ split_key_flat = split_key.reshape(-1)
+
+ # Calculate communication split axis for transposing later
+ if key_is_single_tensor or key_is_mask_like:
+ communication_split = 0
+ else:
+ communication_split = (
+ output_split - (split_key.ndim - 1) if split_key.ndim > 1 else output_split
+ )
+
+ # Step 1: route and send index requests
+
+ sort_idx, send_counts_t, send_displs_t, recv_counts_t, recv_displs_t = (
+ self.__prepare_unordered_comm(split_key_flat, displs)
+ )
+
+ send_counts = send_counts_t.tolist()
+ send_displs = send_displs_t.tolist()
+ recv_counts = recv_counts_t.tolist()
+ recv_displs = recv_displs_t.tolist()
+
+ # Expand counts for multidimensional mask coordinates
+ if key_is_mask_like:
+ mask_dims = len(key)
+ idx_send_counts = [c * mask_dims for c in send_counts]
+ idx_send_displs = [d * mask_dims for d in send_displs]
+ idx_recv_counts = [c * mask_dims for c in recv_counts]
+ idx_recv_displs = [d * mask_dims for d in recv_displs]
+
+ send_indices = torch.stack([k.flatten()[sort_idx] for k in key], dim=1).reshape(-1)
+ recv_indices_flat = torch.zeros(
+ sum(idx_recv_counts), dtype=split_key.dtype, device=self.larray.device
+ )
+ else:
+ idx_send_counts, idx_send_displs = send_counts, send_displs
+ idx_recv_counts, idx_recv_displs = recv_counts, recv_displs
+
+ send_indices = split_key_flat[sort_idx]
+ recv_indices_flat = torch.zeros(
+ sum(idx_recv_counts), dtype=split_key.dtype, device=self.larray.device
+ )
+
+ self.comm.Alltoallv(
+ (send_indices, idx_send_counts, idx_send_displs),
+ (recv_indices_flat, idx_recv_counts, idx_recv_displs),
+ )
+
+ if key_is_mask_like:
+ recv_indices = recv_indices_flat.reshape(sum(recv_counts), len(key))
+ else:
+ recv_indices = recv_indices_flat
+
+ # Step 2: local data lookup based on received indices
+
+ if key_is_mask_like:
+ recv_indices[:, self.split] -= displs[rank]
+ lookup_key = tuple(recv_indices[:, i] for i in range(len(key)))
+ local_vals = self.larray[lookup_key]
+ else:
+ recv_indices -= displs[rank]
+ if key_is_single_tensor:
+ local_vals = self.larray[recv_indices]
+ else:
+ lookup_key = list(key)
+ lookup_key[self.split] = recv_indices
+ local_vals = self.larray[tuple(lookup_key)]
+
+ # Step 3: return data to requesting processes
+
+ # Ensure the indexed elements are aligned along axis 0
+ transpose_axes = list(range(local_vals.ndim))
+ transpose_axes[0], transpose_axes[communication_split] = (
+ transpose_axes[communication_split],
+ transpose_axes[0],
+ )
+ local_vals = local_vals.permute(*transpose_axes)
+
+ feature_shape = list(local_vals.shape[1:])
+ feature_size = 1
+ for dim in feature_shape:
+ feature_size *= dim
+
+ return_send_counts = [c * feature_size for c in recv_counts]
+ return_send_displs = [d * feature_size for d in recv_displs]
+ return_recv_counts = [c * feature_size for c in send_counts]
+ return_recv_displs = [d * feature_size for d in send_displs]
+
+ send_vals = local_vals.reshape(-1)
+ recv_vals_flat = torch.empty(
+ sum(return_recv_counts), dtype=self.larray.dtype, device=self.larray.device
+ )
+
+ self.comm.Alltoallv(
+ (send_vals, return_send_counts, return_send_displs),
+ (recv_vals_flat, return_recv_counts, return_recv_displs),
+ )
+
+ # Step 4: reshape received values and reorder to match original key order
+
+ recv_vals = recv_vals_flat.reshape(-1, *feature_shape)
+
+ # Reverse the sorting applied in Step 1
+ inv_sort_idx = torch.empty_like(sort_idx)
+ inv_sort_idx[sort_idx] = torch.arange(sort_idx.numel(), device=sort_idx.device)
+ unsorted_vals = recv_vals[inv_sort_idx]
+
+ # Restore original dimension order
+ final_vals = unsorted_vals.permute(*transpose_axes)
+
+ # Reshape to match the global output shape expectation
+ if communication_split != output_split:
+ original_local_shape = (
+ output_shape[:communication_split]
+ + split_key.shape
+ + output_shape[output_split + 1 :]
+ )
+ final_vals = final_vals.reshape(original_local_shape)
+
+ indexed_arr = DNDarray(
+ final_vals,
+ gshape=output_shape,
+ dtype=self.dtype,
+ split=output_split,
+ device=self.device,
+ comm=self.comm,
+ balanced=out_is_balanced,
+ )
+
+ if self.ndim > 0:
+ return self.transpose(backwards_transpose_axes), indexed_arr
+ return self, indexed_arr
+
+ def __prepare_unordered_comm(self, split_key_flat: torch.Tensor, displs: tuple) -> tuple:
+ """
+ Helper function for distributed unordered indexing.
+ Determines destination ranks, sorts the key, and computes Alltoallv parameters.
+ """
+ displs_t = torch.tensor(displs, device=self.device.torch_device)
+
+ # map global indices to destination ranks
+ dest_ranks = torch.searchsorted(displs_t[1:], split_key_flat, right=True).to(torch.int64)
+
+ # sort by destination rank to pack memory contiguously
+ sort_idx = torch.argsort(dest_ranks)
+ dest_ranks_sorted = dest_ranks[sort_idx]
+
+ # calculate send_counts and send_displs
+ send_counts = torch.bincount(dest_ranks_sorted, minlength=self.comm.size).to(torch.int64)
+ send_displs = torch.zeros_like(send_counts)
+ send_displs[1:] = torch.cumsum(send_counts, dim=0)[:-1]
+
+ # compose communication matrix, i.e. share `send_counts` information with all processes
+ comm_matrix = torch.zeros(
+ (self.comm.size, self.comm.size),
+ dtype=torch.int64,
+ device=self.device.torch_device,
+ )
+ self.comm.Allgather(send_counts, comm_matrix)
+
+ # comm_matrix columns contain recv_counts for each process
+ recv_counts = comm_matrix[:, self.comm.rank].squeeze(0)
+ recv_displs = torch.zeros_like(recv_counts)
+ recv_displs[1:] = recv_counts.cumsum(0)[:-1]
+
+ return (
+ sort_idx,
+ send_counts,
+ send_displs,
+ recv_counts,
+ recv_displs,
+ )
+
+ def __getitem__(self, key: Indexer) -> DNDarray:
"""
Global getter function for DNDarrays.
- Returns a new DNDarray composed of the elements of the original tensor selected by the indices
- given. This does *NOT* redistribute or rebalance the resulting tensor. If the selection of values is
- unbalanced then the resultant tensor is also unbalanced!
- To redistributed the ``DNDarray`` use :func:`balance()` (issue #187)
+
+ Returns a new DNDarray corresponding to the selection of values from the original DNDarray
+ as specified by `key`. The `key` can be a variety of indexers, including integers, slices,
+ lists, boolean masks, DNDarrays, ndarrays, torch tensors, and a combination thereof.
+
+ The function determines the appropriate method to retrieve the requested data based on the
+ type and structure of `key`, executing MPI communication if the indexing pattern requires
+ data from multiple processes.
+
+ Notes
+ -----
+ The returned DNDarray will have its shape, split, and balanced status determined according
+ to the indexing operation performed. For more details on supported indexing behaviors, see
+ the :doc:`indexing documentation `.
Parameters
----------
- key : int, slice, Tuple[int,...], List[int,...]
- Indices to get from the tensor.
+ key : array-like indexer
+ Indices to get from the ``DNDarray``.
Examples
--------
@@ -908,232 +2309,37 @@ def __getitem__(self, key: Union[int, Tuple[int, ...], List[int, ...]]) -> DNDar
(1/2) >>> tensor([0.])
(2/2) >>> tensor([0., 0.])
"""
- key = getattr(key, "copy()", key)
- l_dtype = self.dtype.torch_type()
- advanced_ind = False
- if isinstance(key, DNDarray) and key.ndim == self.ndim:
- """if the key is a DNDarray and it has as many dimensions as self, then each of the
- entries in the 0th dim refer to a single element. To handle this, the key is split
- into the torch tensors for each dimension. This signals that advanced indexing is
- to be used."""
- # NOTE: this gathers the entire key on every process!!
- # TODO: remove this resplit!!
- key = manipulations.resplit(key)
- if key.larray.dtype in [torch.bool, torch.uint8]:
- key = indexing.nonzero(key)
-
- if key.ndim > 1:
- key = list(key.larray.split(1, dim=1))
- # key is now a list of tensors with dimensions (key.ndim, 1)
- # squeeze singleton dimension:
- key = [key[i].squeeze_(1) for i in range(len(key))]
- else:
- key = [key]
- advanced_ind = True
- elif not isinstance(key, tuple):
- """this loop handles all other cases. DNDarrays which make it to here refer to
- advanced indexing slices, as do the torch tensors. Both DNDaarrys and torch.Tensors
- are cast into lists here by PyTorch. lists mean advanced indexing will be used"""
- h = [slice(None, None, None)] * max(self.ndim, 1)
- if isinstance(key, DNDarray):
- key = manipulations.resplit(key)
- if key.larray.dtype in [torch.bool, torch.uint8]:
- h[0] = torch.nonzero(key.larray).flatten() # .tolist()
- else:
- h[0] = key.larray.tolist()
- elif isinstance(key, torch.Tensor):
- if key.dtype in [torch.bool, torch.uint8]:
- # (coquelin77) i am not certain why this works without being a list. but it works...for now
- h[0] = torch.nonzero(key).flatten() # .tolist()
- else:
- h[0] = key.tolist()
- else:
- h[0] = key
-
- key = list(h)
-
- if isinstance(key, (list, tuple)):
- key = list(key)
- for i, k in enumerate(key):
- # this might be a good place to check if the dtype is there
- try:
- k = manipulations.resplit(k)
- key[i] = k.larray
- except AttributeError:
- pass
-
- # ellipsis
- key = list(key)
- key_classes = [type(n) for n in key]
- # if any(isinstance(n, ellipsis) for n in key):
- n_elips = key_classes.count(type(...))
- if n_elips > 1:
- raise ValueError("key can only contain 1 ellipsis")
- elif n_elips == 1:
- # get which item is the ellipsis
- ell_ind = key_classes.index(type(...))
- kst = key[:ell_ind]
- kend = key[ell_ind + 1 :]
- slices = [slice(None)] * (self.ndim - (len(kst) + len(kend)))
- key = kst + slices + kend
- else:
- key = key + [slice(None)] * (self.ndim - len(key))
-
- self_proxy = self.__torch_proxy__()
- for i in range(len(key)):
- if self.__key_adds_dimension(key, i, self_proxy):
- key[i] = slice(None)
- return self.expand_dims(i)[tuple(key)]
-
- key = tuple(key)
- # assess final global shape
- gout_full = list(self_proxy[key].shape)
-
- # calculate new split axis
- new_split = self.split
- # when slicing, squeezed singleton dimensions may affect new split axis
- if self.split is not None and len(gout_full) < self.ndim:
- if advanced_ind:
- new_split = 0
- else:
- for i in range(len(key[: self.split + 1])):
- if self.__key_is_singular(key, i, self_proxy):
- new_split = None if i == self.split else new_split - 1
-
- key = tuple(key)
- if not self.is_distributed():
- arr = self.__array[key].reshape(gout_full)
- return DNDarray(
- arr, tuple(gout_full), self.dtype, new_split, self.device, self.comm, self.balanced
- )
-
- # else: (DNDarray is distributed)
- arr = torch.tensor([], dtype=self.__array.dtype, device=self.__array.device)
- rank = self.comm.rank
- counts, chunk_starts = self.counts_displs()
- counts, chunk_starts = torch.tensor(counts), torch.tensor(chunk_starts)
- chunk_ends = chunk_starts + counts
- chunk_start = chunk_starts[rank]
- chunk_end = chunk_ends[rank]
-
- if len(key) == 0: # handle empty list
- # this will return an array of shape (0, ...)
- arr = self.__array[key]
-
- """ At the end of the following if/elif/elif block the output array will be set.
- each block handles the case where the element of the key along the split axis
- is a different type and converts the key from global indices to local indices. """
- lout = gout_full.copy()
-
+ if key is None:
+ return self.expand_dims(0)
if (
- isinstance(key[self.split], (list, torch.Tensor, DNDarray, np.ndarray))
- and len(key[self.split]) > 1
+ key is ...
+ or (isinstance(key, slice) and key == slice(None))
+ or (isinstance(key, tuple) and key == ())
):
- # advanced indexing, elements in the split dimension are adjusted to the local indices
- lkey = list(key)
- if isinstance(key[self.split], DNDarray):
- lkey[self.split] = key[self.split].larray
-
- if not isinstance(lkey[self.split], torch.Tensor):
- inds = torch.tensor(
- lkey[self.split], dtype=torch.long, device=self.device.torch_device
- )
- elif lkey[self.split].dtype in [torch.bool, torch.uint8]: # or torch.byte?
- # need to convert the bools to indices
- inds = torch.nonzero(lkey[self.split])
- else:
- inds = lkey[self.split]
- # todo: remove where in favor of nonzero? might be a speed upgrade. testing required
- loc_inds = torch.where((inds >= chunk_start) & (inds < chunk_end))
- # if there are no local indices on a process, then `arr` is empty
- # if local indices exist:
- if len(loc_inds[0]) != 0:
- # select same local indices for other (non-split) dimensions if necessary
- for i, k in enumerate(lkey):
- if isinstance(k, (list, torch.Tensor, DNDarray)) and i != self.split:
- lkey[i] = k[loc_inds]
- # correct local indices for offset
- inds = inds[loc_inds] - chunk_start
- lkey[self.split] = inds
- lout[new_split] = len(inds)
- arr = self.__array[tuple(lkey)].reshape(tuple(lout))
- elif len(loc_inds[0]) == 0:
- if new_split is not None:
- lout[new_split] = len(loc_inds[0])
- else:
- lout = [0] * len(gout_full)
- arr = torch.tensor([], dtype=self.larray.dtype, device=self.larray.device).reshape(
- tuple(lout)
- )
-
- elif isinstance(key[self.split], slice):
- # standard slicing along the split axis,
- # adjust the slice start, stop, and step, then run it on the processes which have the requested data
- key = list(key)
- key[self.split] = stride_tricks.sanitize_slice(key[self.split], self.gshape[self.split])
- key_start, key_stop, key_step = (
- key[self.split].start,
- key[self.split].stop,
- key[self.split].step,
- )
- og_key_start = key_start
- st_pr = torch.where(key_start < chunk_ends)[0]
- st_pr = st_pr[0] if len(st_pr) > 0 else self.comm.size
- sp_pr = torch.where(key_stop >= chunk_starts)[0]
- sp_pr = sp_pr[-1] if len(sp_pr) > 0 else 0
- actives = list(range(st_pr, sp_pr + 1))
- if rank in actives:
- key_start = 0 if rank != actives[0] else key_start - chunk_starts[rank]
- key_stop = counts[rank] if rank != actives[-1] else key_stop - chunk_starts[rank]
- key_start, key_stop = self.__xitem_get_key_start_stop(
- rank, actives, key_start, key_stop, key_step, chunk_ends, og_key_start
- )
- key[self.split] = slice(key_start, key_stop, key_step)
- lout[new_split] = (
- math.ceil((key_stop - key_start) / key_step)
- if key_step is not None
- else key_stop - key_start
- )
- arr = self.__array[tuple(key)].reshape(lout)
- else:
- lout[new_split] = 0
- arr = torch.empty(lout, dtype=self.__array.dtype, device=self.__array.device)
-
- elif self.__key_is_singular(key, self.split, self_proxy):
- # getting one item along split axis:
- key = list(key)
- if isinstance(key[self.split], list):
- key[self.split] = key[self.split].pop()
- elif isinstance(key[self.split], (torch.Tensor, DNDarray, np.ndarray)):
- key[self.split] = key[self.split].item()
- # translate negative index
- if key[self.split] < 0:
- key[self.split] += self.gshape[self.split]
-
- active_rank = torch.where(key[self.split] >= chunk_starts)[0][-1].item()
- # slice `self` on `active_rank`, allocate `arr` on all other ranks in preparation for Bcast
- if rank == active_rank:
- key[self.split] -= chunk_start.item()
- arr = self.__array[tuple(key)].reshape(tuple(lout))
- else:
- arr = torch.empty(tuple(lout), dtype=self.larray.dtype, device=self.larray.device)
- # broadcast result
- # TODO: Replace with `self.comm.Bcast(arr, root=active_rank)` after fixing #784
- arr = self.comm.bcast(arr, root=active_rank)
- if arr.device != self.larray.device:
- # todo: remove when unnecessary (also after #784)
- arr = arr.to(device=self.larray.device)
+ return self
- return DNDarray(
- arr.type(l_dtype),
- gout_full if isinstance(gout_full, tuple) else tuple(gout_full),
- self.dtype,
- new_split,
- self.device,
- self.comm,
- balanced=True if new_split is None else None,
+ # key processing returns a ProcessedKey namedtuple
+ self, processed_key = _resolve_indexing_state(
+ self, key, return_local_indices=True, op="get"
)
+ # dispatch to appropriate getitem method
+ op = processed_key.op_type
+ # print("DEBUGGING: Operation type:", op)
+
+ if op == "scalar":
+ return self.__getitem_scalar(processed_key)
+ elif op == "distr_mask":
+ return self.__getitem_mask(processed_key, key)
+ elif op == "distributed":
+ return self.__getitem_advanced_distributed(processed_key)
+ elif op == "slice":
+ return self.__getitem_slice(processed_key)
+ elif op == "descending_slice":
+ return self.__getitem_descending_slice_distributed(processed_key)
+ elif op in ("local_mask", "advanced"):
+ return self.__getitem_advanced_local(processed_key, key)
+
if torch.cuda.device_count() > 0:
def gpu(self) -> DNDarray:
@@ -1183,18 +2389,6 @@ def is_distributed(self) -> bool:
"""
return self.split is not None and self.comm.is_distributed()
- @staticmethod
- def __key_is_singular(key: any, axis: int, self_proxy: torch.Tensor) -> bool:
- # determine if the key gets a singular item
- zeros = (0,) * (self_proxy.ndim - 1)
- return self_proxy[(*zeros[:axis], key[axis], *zeros[axis:])].ndim == 0
-
- @staticmethod
- def __key_adds_dimension(key: any, axis: int, self_proxy: torch.Tensor) -> bool:
- # determine if the key adds a new dimension
- zeros = (0,) * (self_proxy.ndim - 1)
- return self_proxy[(*zeros[:axis], key[axis], *zeros[axis:])].ndim == 2
-
def item(self):
"""
Returns the only element of a 1-element :class:`DNDarray`.
@@ -1218,9 +2412,13 @@ def __len__(self) -> int:
"""
The length of the ``DNDarray``, i.e. the number of items in the first dimension.
"""
- return self.shape[0]
+ try:
+ len = self.shape[0]
+ return len
+ except IndexError:
+ raise TypeError("len() of unsized DNDarray")
- def numpy(self) -> np.array:
+ def numpy(self) -> np.typing.NDArray[Any]:
"""
Returns a copy of the :class:`DNDarray` as numpy ndarray. If the ``DNDarray`` resides on the GPU, the underlying data will be copied to the CPU first.
@@ -1250,7 +2448,7 @@ def __repr__(self) -> str:
"""
return printing.__repr__(self)
- def ravel(self):
+ def ravel(self) -> "DNDarray":
"""
Flattens the ``DNDarray``.
@@ -1269,8 +2467,8 @@ def ravel(self):
return manipulations.ravel(self)
def redistribute_(
- self, lshape_map: Optional[torch.Tensor] = None, target_map: Optional[torch.Tensor] = None
- ):
+ self, lshape_map: torch.Tensor | None = None, target_map: torch.Tensor | None = None
+ ) -> None:
"""
Redistributes the data of the :class:`DNDarray` *along the split axis* to match the given target map.
This function does not modify the non-split dimensions of the ``DNDarray``.
@@ -1422,9 +2620,9 @@ def redistribute_(
def __redistribute_shuffle(
self,
- snd_pr: Union[int, torch.Tensor],
- send_amt: Union[int, torch.Tensor],
- rcv_pr: Union[int, torch.Tensor],
+ snd_pr: int | torch.Tensor,
+ send_amt: int | torch.Tensor,
+ rcv_pr: int | torch.Tensor,
snd_dtype: torch.dtype,
):
"""
@@ -1514,16 +2712,14 @@ def resplit_(self, axis: int = None):
# sanitize the axis to check whether it is in range
axis = sanitize_axis(self.shape, axis)
+ self.__partitions_dict__ = None
+
# early out for unchanged content
if self.comm.size == 1:
self.__split = axis
- if axis is None:
- self.__partitions_dict__ = None
if axis == self.split:
return self
- self.__partitions_dict__ = None
-
if axis is None:
gathered = torch.empty(
self.shape, dtype=self.dtype.torch_type(), device=self.device.torch_device
@@ -1567,322 +2763,560 @@ def resplit_(self, axis: int = None):
return self
- def __setitem__(
- self,
- key: Union[int, Tuple[int, ...], List[int, ...]],
- value: Union[float, DNDarray, torch.Tensor],
- ):
+ def __setitem_scalar(self, p: ProcessedKey, value: "DNDarray", value_is_scalar: bool) -> None:
+ if p.root is not None:
+ if self.comm.rank == p.root:
+ indexed_proxy = self.__torch_proxy__()[p.key]
+ if indexed_proxy.names.count("split") != 0:
+ indexed_lshape_map = self.lshape_map[:, 1:]
+ if value.lshape_map != indexed_lshape_map:
+ try:
+ value.redistribute_(target_map=indexed_lshape_map)
+ except ValueError:
+ raise ValueError(
+ f"cannot assign value to indexed DNDarray because "
+ f"distribution schemes do not match: "
+ f"{value.lshape_map} vs. {indexed_lshape_map}"
+ )
+ self.__set(p.key, value)
+ else:
+ if not value_is_scalar:
+ value = sanitation.sanitize_distribution(value, target=self[p.key])
+ self.__set(p.key, value)
+
+ def __setitem_slice(self, p: ProcessedKey, value: "DNDarray", value_is_scalar: bool) -> None:
+ """
+ Assigns a value array using standard slicing. If `value` is distributed, it might be redistributed to align with the target slice before assignment.
"""
- Global item setter
+ if not self.is_distributed() and not value.is_distributed():
+ self.__set(p.key, value)
+ return
- Parameters
- ----------
- key : Union[int, Tuple[int,...], List[int,...]]
- Index/indices to be set
- value: Union[float, DNDarray,torch.Tensor]
- Value to be set to the specified positions in the DNDarray (self)
+ if self.is_distributed() and not value_is_scalar:
+ if not value.is_distributed():
+ value = factories.array(
+ value.larray,
+ dtype=value.dtype,
+ split=p.output_split,
+ device=self.device,
+ comm=self.comm,
+ )
+ else:
+ if value.split != p.output_split:
+ raise RuntimeError(
+ f"Cannot assign distributed `value` with split axis {value.split} "
+ f"to indexed DNDarray with split axis {p.output_split}."
+ )
+ target_shape = torch.tensor(
+ tuple(self.larray[p.key].shape), device=self.device.torch_device
+ )
+ target_map = torch.zeros(
+ (self.comm.size, len(target_shape)),
+ dtype=torch.int64,
+ device=self.device.torch_device,
+ )
+ self.comm.Allgather(target_shape, target_map)
+ value.redistribute_(target_map=target_map)
- Notes
- -----
- If a ``DNDarray`` is given as the value to be set then the split axes are assumed to be equal.
- If they are not, PyTorch will raise an error when the values are attempted to be set on the local array
+ self.__set(p.key, value)
- Examples
- --------
- >>> a = ht.zeros((4, 5), split=0)
- (1/2) >>> tensor([[0., 0., 0., 0., 0.],
- [0., 0., 0., 0., 0.]])
- (2/2) >>> tensor([[0., 0., 0., 0., 0.],
- [0., 0., 0., 0., 0.]])
- >>> a[1:4, 1] = 1
- >>> a
- (1/2) >>> tensor([[0., 0., 0., 0., 0.],
- [0., 1., 0., 0., 0.]])
- (2/2) >>> tensor([[0., 1., 0., 0., 0.],
- [0., 1., 0., 0., 0.]])
+ def __setitem_advanced_local(
+ self, p: ProcessedKey, original_key, value: "DNDarray", value_is_scalar: bool
+ ) -> None:
"""
- key = getattr(key, "copy()", key)
- try:
- if value.split != self.split:
- val_split = int(value.split)
- sp = self.split
- warnings.warn(
- f"\nvalue.split {val_split} not equal to this DNDarray's split:"
- f" {sp}. this may cause errors or unwanted behavior",
- category=RuntimeWarning,
- )
- except (AttributeError, TypeError):
- pass
-
- # NOTE: for whatever reason, there is an inplace op which interferes with the abstraction
- # of this next block of code. this is shared with __getitem__. I attempted to abstract it
- # in a standard way, but it was causing errors in the test suite. If someone else is
- # motived to do this they are welcome to, but i have no time right now
- # print(key)
- if isinstance(key, DNDarray) and key.ndim == self.ndim:
- """if the key is a DNDarray and it has as many dimensions as self, then each of the
- entries in the 0th dim refer to a single element. To handle this, the key is split
- into the torch tensors for each dimension. This signals that advanced indexing is
- to be used."""
- key = manipulations.resplit(key)
- if key.larray.dtype in [torch.bool, torch.uint8]:
- key = indexing.nonzero(key)
-
- if key.ndim > 1:
- key = list(key.larray.split(1, dim=1))
- # key is now a list of tensors with dimensions (key.ndim, 1)
- # squeeze singleton dimension:
- key = [key[i].squeeze_(1) for i in range(len(key))]
+ Handles local advanced indexing assignments.
+ """
+ self.__setitem_slice(p, value, value_is_scalar)
+
+ def __setitem_descending_slice_distributed(
+ self, p: ProcessedKey, value: "DNDarray", value_is_scalar: bool
+ ) -> None:
+ """
+ Handles assignment via negative-step slicing. Flips the `value` array and redistributes
+ it to align with the descending split key before performing the local assignment.
+ """
+ flipped_value = manipulations.flip(value, axis=p.output_split)
+ split_key = factories.array(
+ p.key[self.split], is_split=0, device=self.device, comm=self.comm
+ )
+ if not flipped_value.is_distributed():
+ flipped_value = factories.array(
+ flipped_value.larray,
+ dtype=flipped_value.dtype,
+ split=p.output_split,
+ device=self.device,
+ comm=self.comm,
+ )
+ target_map = flipped_value.lshape_map
+ target_map[:, p.output_split] = split_key.lshape_map[:, 0]
+ flipped_value.redistribute_(target_map=target_map)
+ self.__set(p.key, flipped_value)
+
+ def __setitem_mask(
+ self, p: ProcessedKey, original_key, value: "DNDarray", value_is_scalar: bool
+ ) -> None:
+ """
+ Handles assignment using boolean masks. If `value` is distributed, it will be redistributed to match the number of True elements in the local mask before assignment. If `value` is not distributed, it will be assigned directly to the masked positions on each process, with PyTorch handling any necessary broadcasting.
+ """
+ pytorch_key = p.key
+
+ if isinstance(pytorch_key, tuple):
+ for k in pytorch_key:
+ if isinstance(k, torch.Tensor) and k.dtype in (torch.bool, torch.uint8):
+ local_mask = k
+ break
+ else:
+ local_mask = pytorch_key
+
+ if value_is_scalar:
+ if hasattr(value, "larray"):
+ scalar_torch = value.larray
+ else:
+ scalar_torch = torch.as_tensor(value, device=self.device.torch_device)
+ scalar_torch = scalar_torch.type(self.dtype.torch_type())
+ self.larray[pytorch_key] = scalar_torch
+ else:
+ if isinstance(value, DNDarray) and value.is_distributed():
+ expected_elements = int(local_mask.sum().item())
+ if value.lshape[0] != expected_elements:
+ raise ValueError(
+ f"Shape mismatch: Cannot assign distributed array with local shape {value.lshape} "
+ f"to a mask requiring {expected_elements} elements on rank {self.comm.rank}."
+ )
+
+ # value perfectly aligns
+ value_torch = value.larray
+ self.larray[pytorch_key] = value_torch.type(self.dtype.torch_type())
+
else:
- key = [key]
- elif not isinstance(key, tuple):
- """this loop handles all other cases. DNDarrays which make it to here refer to
- advanced indexing slices, as do the torch tensors. Both DNDaarrys and torch.Tensors
- are cast into lists here by PyTorch. lists mean advanced indexing will be used"""
- h = [slice(None, None, None)] * self.ndim
- if isinstance(key, DNDarray):
- key = manipulations.resplit(key)
- if key.larray.dtype in [torch.bool, torch.uint8]:
- h[0] = torch.nonzero(key.larray).flatten() # .tolist()
+ # Value is a non-distributed array -> MPI prefix sum needed
+ if hasattr(value, "larray"):
+ value_torch = value.larray
else:
- h[0] = key.larray.tolist()
- elif isinstance(key, torch.Tensor):
- if key.dtype in [torch.bool, torch.uint8]:
- # (coquelin77) im not sure why this works without being a list...but it does...for now
- h[0] = torch.nonzero(key).flatten() # .tolist()
+ value_torch = torch.as_tensor(value, device=self.device.torch_device)
+
+ # distinguish between exact-shape masks and 1D row-filtering masks
+ is_row_mask = local_mask.ndim == 1 and self.ndim > 1
+
+ if not is_row_mask and value_torch.ndim == 1:
+ # N-D mask on N-D array -> flattens into 1D sequence, requires MPI prefix sum
+ local_mask_flat = local_mask.flatten()
+ local_true = int(local_mask_flat.sum().item())
+
+ if self.comm.size > 1:
+ if self.comm.rank == 0:
+ offset = 0
+ _ = self.comm.exscan(local_true)
+ else:
+ offset = self.comm.exscan(local_true)
+ else:
+ offset = 0
+
+ rhs_local = value_torch[offset : offset + local_true].type(
+ self.dtype.torch_type()
+ )
+
+ x_flat = self.larray.view(-1)
+ x_flat[local_mask_flat] = rhs_local
else:
- h[0] = key.tolist()
- else:
- h[0] = key
- key = list(h)
+ # PyTorch assigns and broadcasts natively
+ self.larray[pytorch_key] = value_torch.type(self.dtype.torch_type())
- # key must be torch-proof
- if isinstance(key, (list, tuple)):
- key = list(key)
- for i, k in enumerate(key):
- try: # extract torch tensor
- k = manipulations.resplit(k)
- key[i] = k.larray
- except AttributeError:
- pass
- # remove bools from a torch tensor in favor of indexes
- try:
- if key[i].dtype in [torch.bool, torch.uint8]:
- key[i] = torch.nonzero(key[i]).flatten()
- except (AttributeError, TypeError):
- pass
+ def __setitem_advanced_distributed(
+ self, p: ProcessedKey, original_key, value: "DNDarray", value_is_scalar: bool
+ ) -> None:
+ """
+ Handles advanced indexing assignments where the indexing key is distributed. This method ensures that the value array is properly aligned and redistributed if necessary before performing the local assignment on each process.
+ """
+ # check distribution status of the indexing key
+ split_key_orig = (
+ original_key[self.split] if isinstance(original_key, tuple) else original_key
+ )
+ key_is_distributed = (
+ isinstance(split_key_orig, DNDarray) and split_key_orig.is_distributed()
+ )
- key = list(key)
+ # reject implicit cross-distribution assignments
+ if key_is_distributed and not value.is_distributed() and not value_is_scalar:
+ raise ValueError(
+ f"Distribution mismatch: index distributed={key_is_distributed}, value distributed={value.is_distributed()}. "
+ "Cannot assign a non-distributed value array using a distributed index. "
+ "Please distribute the value array or use a non-distributed index."
+ )
- # ellipsis stuff
- key_classes = [type(n) for n in key]
- # if any(isinstance(n, ellipsis) for n in key):
- n_elips = key_classes.count(type(...))
- if n_elips > 1:
- raise ValueError("key can only contain 1 ellipsis")
- elif n_elips == 1:
- # get which item is the ellipsis
- ell_ind = key_classes.index(type(...))
- kst = key[:ell_ind]
- kend = key[ell_ind + 1 :]
- slices = [slice(None)] * (self.ndim - (len(kst) + len(kend)))
- key = kst + slices + kend
- # ---------- end ellipsis stuff -------------
-
- for c, k in enumerate(key):
- try:
- key[c] = k.item()
- except (AttributeError, ValueError, RuntimeError):
- pass
+ if value.is_distributed():
+ self.__setitem_unordered(
+ key=p.key,
+ key_is_mask_like=p.key_is_mask_like,
+ value=value,
+ key_is_single_tensor=isinstance(p.key, torch.Tensor),
+ counts=self.counts_displs()[0],
+ displs=self.counts_displs()[1],
+ rank=self.comm.rank,
+ backwards_transpose_axes=p.backwards_transpose_axes,
+ )
+ return
+ counts, displs = self.counts_displs()
rank = self.comm.rank
- if self.split is not None:
- counts, chunk_starts = self.counts_displs()
- else:
- counts, chunk_starts = 0, [0] * self.comm.size
- counts = torch.tensor(counts, device=self.device.torch_device)
- chunk_starts = torch.tensor(chunk_starts, device=self.device.torch_device)
- chunk_ends = chunk_starts + counts
- chunk_start = chunk_starts[rank]
- chunk_end = chunk_ends[rank]
- # determine which elements are on the local process (if the key is a torch tensor)
- try:
- # if isinstance(key[self.split], torch.Tensor):
- filter_key = torch.nonzero(
- (chunk_start <= key[self.split]) & (key[self.split] < chunk_end)
+ key_is_single_tensor = isinstance(p.key, torch.Tensor)
+
+ if (
+ value_is_scalar
+ and isinstance(original_key, tuple)
+ and len(original_key) == self.ndim
+ and all(
+ isinstance(k, DNDarray) and k.ndim == 1 and k.dtype in (types.int32, types.int64)
+ for k in original_key
)
- for k in range(len(key)):
- try:
- key[k] = key[k][filter_key].flatten()
- except TypeError:
- pass
- except TypeError: # this will happen if the key doesnt have that many
- pass
+ ):
+ global_indices = []
+ for k in original_key:
+ k_full = k.copy()
+ k_full.resplit_(None)
+ global_indices.append(k_full.larray)
+
+ idx_split_global = global_indices[self.split]
+ local_offset = displs[rank]
+ local_size = counts[rank]
+
+ mask = (idx_split_global >= local_offset) & (
+ idx_split_global < local_offset + local_size
+ )
+ if not mask.any():
+ return
- key = tuple(key)
+ lhs_index = []
+ for dim, gind in enumerate(global_indices):
+ sel = gind[mask]
+ if dim == self.split:
+ sel = sel - local_offset
+ lhs_index.append(sel)
+ lhs_index = tuple(lhs_index)
- if not self.is_distributed():
- return self.__setter(key, value) # returns None
+ if hasattr(value, "larray"):
+ scalar_torch = value.larray
+ else:
+ scalar_torch = torch.as_tensor(value, device=self.device.torch_device)
+ scalar_torch = scalar_torch.type(self.dtype.torch_type())
- # raise RuntimeError("split axis of array and the target value are not equal") removed
- # this will occur if the local shapes do not match
- rank = self.comm.rank
- ends = []
- for pr in range(self.comm.size):
- _, _, e = self.comm.chunk(self.shape, self.split, rank=pr)
- ends.append(e[self.split].stop - e[self.split].start)
- ends = torch.tensor(ends, device=self.device.torch_device)
- chunk_ends = ends.cumsum(dim=0)
- chunk_starts = torch.tensor([0] + chunk_ends.tolist(), device=self.device.torch_device)
- _, _, chunk_slice = self.comm.chunk(self.shape, self.split)
- chunk_start = chunk_slice[self.split].start
- chunk_end = chunk_slice[self.split].stop
-
- self_proxy = self.__torch_proxy__()
-
- # if the value is a DNDarray, the divisions need to be balanced:
- # this means that we need to know how much data is where for both DNDarrays
- # if the value data is not in the right place, then it will need to be moved
-
- if isinstance(key[self.split], slice):
- key = list(key)
- key_start = key[self.split].start if key[self.split].start is not None else 0
- key_stop = (
- key[self.split].stop
- if key[self.split].stop is not None
- else self.gshape[self.split]
- )
- if key_stop < 0:
- key_stop = self.gshape[self.split] + key[self.split].stop
- key_step = key[self.split].step
- og_key_start = key_start
- st_pr = torch.where(key_start < chunk_ends)[0]
- st_pr = st_pr[0] if len(st_pr) > 0 else self.comm.size
- sp_pr = torch.where(key_stop >= chunk_starts)[0]
- sp_pr = sp_pr[-1] if len(sp_pr) > 0 else 0
- actives = list(range(st_pr, sp_pr + 1))
-
- if (
- isinstance(value, type(self))
- and value.split is not None
- and value.shape[self.split] != self.shape[self.split]
- ):
- # setting elements in self with a DNDarray which is not the same size in the
- # split dimension
- local_keys = []
- # below is used if the target needs to be reshaped
- target_reshape_map = torch.zeros(
- (self.comm.size, self.ndim), dtype=torch.int64, device=self.device.torch_device
- )
- for r in range(self.comm.size):
- if r not in actives:
- loc_key = key.copy()
- loc_key[self.split] = slice(0, 0, 0)
- else:
- key_start_l = 0 if r != actives[0] else key_start - chunk_starts[r]
- key_stop_l = ends[r] if r != actives[-1] else key_stop - chunk_starts[r]
- key_start_l, key_stop_l = self.__xitem_get_key_start_stop(
- r, actives, key_start_l, key_stop_l, key_step, chunk_ends, og_key_start
- )
- loc_key = key.copy()
- loc_key[self.split] = slice(key_start_l, key_stop_l, key_step)
+ self.larray[lhs_index] = scalar_torch
+ return
- gout_full = torch.tensor(
- self_proxy[tuple(loc_key)].shape, device=self.device.torch_device
- )
- target_reshape_map[r] = gout_full
- local_keys.append(loc_key)
+ if key_is_single_tensor:
+ split_key = p.key
+ split_key_flat = split_key.reshape(-1)
+ local_indices = torch.nonzero(
+ (split_key_flat >= displs[rank]) & (split_key_flat < displs[rank] + counts[rank])
+ ).flatten()
- key = local_keys[rank]
- value = value.redistribute(target_map=target_reshape_map)
+ if local_indices.numel() > 0:
+ key_local = split_key_flat[local_indices] - displs[rank]
- if rank not in actives:
- return # non-active ranks can exit here
+ if value_is_scalar:
+ rhs = value.larray.type(self.dtype.torch_type())
+ else:
+ # flatten leading dimensions of value.larray that correspond to the multi-dimensional key
+ rhs_view = value.larray.reshape(-1, *value.larray.shape[split_key.ndim :])
+ rhs = rhs_view[local_indices].type(self.dtype.torch_type())
- chunk_starts_v = target_reshape_map[:, self.split]
- value_slice = [slice(None, None, None)] * value.ndim
- step2 = key_step if key_step is not None else 1
- key_start = (chunk_starts_v[rank] - og_key_start).item()
+ if self.larray.is_cuda:
+ key_local, rhs = _resolve_duplicate_indices(key_local, rhs, self.larray.shape)
- key_start = max(key_start, 0)
- key_stop = key_start + key_stop
- slice_loc = min(self.split, value.ndim - 1)
- value_slice[slice_loc] = slice(
- key_start, math.ceil(torch.true_divide(key_stop, step2)), 1
- )
+ self.larray[key_local] = rhs
+ return
- self.__setter(tuple(key), value.larray)
- return
+ if isinstance(original_key, tuple):
+ original_split_axis = p.backwards_transpose_axes[self.split]
+ raw_split_part = original_key[original_split_axis]
+ else:
+ raw_split_part = original_key
+
+ if isinstance(raw_split_part, DNDarray):
+ split_key = raw_split_part.larray
+ elif isinstance(raw_split_part, torch.Tensor):
+ split_key = raw_split_part
+ else:
+ split_key = p.key[self.split]
+
+ if isinstance(split_key, DNDarray):
+ split_key = split_key.larray
+
+ if split_key.dtype == torch.bool:
+ split_key = torch.nonzero(split_key, as_tuple=False).flatten()
- # if rank in actives:
- if rank not in actives:
- return # non-active ranks can exit here
- key_start = 0 if rank != actives[0] else key_start - chunk_starts[rank]
- key_stop = ends[rank] if rank != actives[-1] else key_stop - chunk_starts[rank]
- key_start, key_stop = self.__xitem_get_key_start_stop(
- rank, actives, key_start, key_stop, key_step, chunk_ends, og_key_start
+ local_offset = displs[rank]
+ local_size = counts[rank]
+
+ if hasattr(value, "larray"):
+ value_torch = value.larray
+ else:
+ value_torch = torch.as_tensor(value, device=self.device.torch_device)
+
+ feature_dims = self.larray.ndim - (self.split + 1)
+
+ if value_is_scalar:
+ value_key_start_dim = 0
+ else:
+ value_key_start_dim = value_torch.ndim - split_key.ndim - feature_dims
+ if value_key_start_dim < 0:
+ raise RuntimeError("value_key_start_dim < 0 – inconsistent shapes")
+
+ local_split_axis = self.split
+
+ base_index = [slice(None)] * self.larray.ndim
+ if isinstance(original_key, tuple):
+ for dim, k_part in enumerate(original_key):
+ if dim == self.split:
+ continue
+ if isinstance(k_part, DNDarray):
+ base_index[dim] = k_part.larray
+ else:
+ base_index[dim] = k_part
+
+ self.__advanced_setitem_unordered_local(
+ x_local=self.larray,
+ split_key=split_key,
+ value_torch=value_torch,
+ split_axis=local_split_axis,
+ value_key_start_dim=value_key_start_dim,
+ local_offset=local_offset,
+ local_size=local_size,
+ value_is_scalar=value_is_scalar,
+ out_dtype=self.dtype.torch_type(),
+ base_index=tuple(base_index),
+ )
+
+ def __setitem_unordered(
+ self,
+ key: tuple | list | torch.Tensor,
+ key_is_mask_like: bool,
+ value: "DNDarray",
+ key_is_single_tensor: bool,
+ counts: tuple,
+ displs: tuple,
+ rank: int,
+ backwards_transpose_axes: tuple,
+ ) -> DNDarray:
+ """
+ Handles the MPI communication when assigning a distributed
+ value to a distributed array with unordered global indices.
+ """
+ # distribution of `key` and `value` must be aligned
+ if key_is_mask_like:
+ # redistribute `value` to match distribution of `key` in one pass
+ split_key = key[self.split]
+ global_split_key = factories.array(
+ split_key, is_split=0, device=self.device, comm=self.comm, copy=False
)
- key[self.split] = slice(key_start, key_stop, key_step)
-
- # todo: need to slice the values to be the right size...
- if isinstance(value, (torch.Tensor, type(self))):
- # if its a torch tensor, it is assumed to exist on all processes
- value_slice = [slice(None, None, None)] * value.ndim
- step2 = key_step if key_step is not None else 1
- key_start = (chunk_starts[rank] - og_key_start).item()
- key_start = max(key_start, 0)
- key_stop = key_start + key_stop
- slice_loc = min(self.split, value.ndim - 1)
- value_slice[slice_loc] = slice(
- key_start, math.ceil(torch.true_divide(key_stop, step2)), 1
- )
- self.__setter(tuple(key), value[tuple(value_slice)])
+ target_map = value.lshape_map
+ target_map[:, value.split] = global_split_key.lshape_map[:, 0]
+ value.redistribute_(target_map=target_map)
+ else:
+ # redistribute split-axis `key` to match distribution of `value` in one pass
+ if key_is_single_tensor:
+ # key is a single torch.Tensor
+ split_key = key
else:
- self.__setter(tuple(key), value)
- elif isinstance(key[self.split], (torch.Tensor, list)):
- key = list(key)
- key[self.split] -= chunk_start
- if len(key[self.split]) != 0:
- self.__setter(tuple(key), value)
+ split_key = key[self.split]
+ global_split_key = factories.array(
+ split_key, is_split=0, device=self.device, comm=self.comm, copy=False
+ )
+ target_map = global_split_key.lshape_map
+ target_map[:, 0] = value.lshape_map[:, value.split]
+ global_split_key.redistribute_(target_map=target_map)
+ split_key = global_split_key.larray
+
+ # key and value are now aligned
+
+ # prepare for `value` Alltoallv:
+ # work along axis 0, transpose if necessary
+ transpose_axes = list(range(value.ndim))
+ transpose_axes[0], transpose_axes[value.split] = (
+ transpose_axes[value.split],
+ transpose_axes[0],
+ )
+ value = value.transpose(transpose_axes)
- elif key[self.split] in range(chunk_start, chunk_end):
- key = list(key)
- key[self.split] = key[self.split] - chunk_start
- self.__setter(tuple(key), value)
+ split_key_flat = split_key.reshape(-1)
+ sort_idx, send_counts, send_displs, recv_counts, recv_displs = (
+ self.__prepare_unordered_comm(split_key_flat, displs)
+ )
+
+ # allocate send buffer: add 1 column to store sent indices
+ send_buf_shape = list(value.lshape)
+ if value.ndim < 2:
+ send_buf_shape.append(1)
+ if key_is_mask_like:
+ send_buf_shape[-1] += len(key)
+ else:
+ send_buf_shape[-1] += 1
+ send_buf = torch.zeros(
+ send_buf_shape, dtype=value.dtype.torch_type(), device=self.device.torch_device
+ )
+
+ # pack the send_buf
+ if sort_idx.numel() > 0:
+ if value.ndim < 2:
+ send_buf[:, :-1] = value.larray[sort_idx].unsqueeze(1)
+ else:
+ send_buf[:, :-1] = value.larray[sort_idx]
- elif key[self.split] < 0:
+ if key_is_mask_like:
+ for i in range(-len(key), 0):
+ send_buf[:, i] = key[i + len(key)][sort_idx]
+ else:
+ send_buf[:, -1] = split_key_flat[sort_idx].to(send_buf.dtype)
+
+ # allocate receive buffer, with 1 extra column for incoming indices
+ recv_buf_shape = value.lshape_map[self.comm.rank]
+ recv_buf_shape[value.split] = recv_counts.sum()
+ recv_buf_shape = recv_buf_shape.tolist()
+ if value.ndim < 2:
+ recv_buf_shape.append(1)
+ if key_is_mask_like:
+ recv_buf_shape[-1] += len(key)
+ else:
+ recv_buf_shape[-1] += 1
+ recv_buf_shape = tuple(recv_buf_shape)
+ recv_buf = torch.zeros(
+ recv_buf_shape, dtype=value.dtype.torch_type(), device=self.device.torch_device
+ )
+ # perform Alltoallv along the 0 axis
+ send_counts, send_displs, recv_counts, recv_displs = (
+ send_counts.tolist(),
+ send_displs.tolist(),
+ recv_counts.tolist(),
+ recv_displs.tolist(),
+ )
+ self.comm.Alltoallv(
+ (send_buf, send_counts, send_displs), (recv_buf, recv_counts, recv_displs)
+ )
+ del send_buf
+
+ if key_is_mask_like:
key = list(key)
- if self.gshape[self.split] + key[self.split] in range(chunk_start, chunk_end):
- key[self.split] = key[self.split] + self.shape[self.split] - chunk_start
- self.__setter(tuple(key), value)
+ # extract incoming indices from recv_buf
+ recv_indices = recv_buf[..., -len(key) :]
+ # correct split-axis indices for rank offset
+ recv_indices[:, 0] -= displs[rank]
+ key = recv_indices.split(1, dim=1)
+ key = [key[i].squeeze_(1) for i in range(len(key))]
+ # remove indices from recv_buf
+ recv_buf = recv_buf[..., : -len(key)]
+ else:
+ # store incoming indices in int 1-D tensor and correct for rank offset
+ recv_indices = recv_buf[..., -1].type(torch.int64) - displs[rank]
+ # remove last column from recv_buf
+ recv_buf = recv_buf[..., :-1]
+
+ # replace split-axis key with incoming local indices
+ if key_is_single_tensor:
+ key = recv_indices
+ else:
+ key = list(key)
+ key[self.split] = recv_indices
+ key = tuple(key)
+
+ # transpose back value and recv_buf if necessary, wrap recv_buf in DNDarray
+ value = value.transpose(transpose_axes)
+ if value.ndim < 2:
+ recv_buf.squeeze_(1)
+ recv_buf = DNDarray(
+ recv_buf.permute(*transpose_axes),
+ gshape=value.gshape,
+ dtype=value.dtype,
+ split=value.split,
+ device=value.device,
+ comm=value.comm,
+ balanced=value.balanced,
+ )
+ # set local elements of `self` to corresponding elements of `value`
+ self.__set(key, recv_buf)
+ if self.ndim > 0:
+ return self.transpose(backwards_transpose_axes)
+ return self
- def __setter(
+ def __setitem__(
self,
- key: Union[int, Tuple[int, ...], List[int, ...]],
- value: Union[float, DNDarray, torch.Tensor],
+ key: Indexer,
+ value: float | "DNDarray" | torch.Tensor,
):
"""
- Utility function for checking ``value`` and forwarding to :func:``__setitem__``
+ Global item setter for DNDarrays.
- Raises
- ------
- NotImplementedError
- If the type of ``value`` ist not supported
- """
- if np.isscalar(value):
- self.__array.__setitem__(key, value)
- elif isinstance(value, DNDarray):
- self.__array.__setitem__(key, value.__array)
- elif isinstance(value, torch.Tensor):
- self.__array.__setitem__(key, value.data)
- elif isinstance(value, (list, tuple)):
- value = torch.tensor(value, device=self.device.torch_device)
- self.__array.__setitem__(key, value.data)
- elif isinstance(value, np.ndarray):
- value = torch.from_numpy(value)
- self.__array.__setitem__(key, value.data)
+ Assigns values to the specified positions in the ``DNDarray``. The `key` can be a variety
+ of indexers, including integers, slices, lists, boolean masks, DNDarrays, ndarrays,
+ torch tensors, or a combination thereof.
+
+ If a distributed ``DNDarray`` is given as the `value` to be set, this function will
+ automatically attempt to align its distribution scheme (split axis and local shapes)
+ with the target indexed array via MPI communication. If the distributions cannot be
+ safely aligned, a ``ValueError`` or ``RuntimeError`` is raised.
+
+ Parameters
+ ----------
+ key : array-like indexer
+ Index/indices to be set
+ value: float | "DNDarray" | torch.Tensor
+ Value to be set to the specified positions in the DNDarray (self)
+
+ Notes
+ -----
+ For more details on supported indexing behaviors, see the :doc:`indexing documentation `.
+
+ Examples
+ --------
+ >>> a = ht.zeros((4, 5), split=0)
+ (1/2) >>> tensor([[0., 0., 0., 0., 0.],
+ [0., 0., 0., 0., 0.]])
+ (2/2) >>> tensor([[0., 0., 0., 0., 0.],
+ [0., 0., 0., 0., 0.]])
+ >>> a[1:4, 1] = 1
+ >>> a
+ (1/2) >>> tensor([[0., 0., 0., 0., 0.],
+ [0., 1., 0., 0., 0.]])
+ (2/2) >>> tensor([[0., 1., 0., 0., 0.],
+ [0., 1., 0., 0., 0.]])
+ """
+ try:
+ value = factories.array(value)
+ except TypeError:
+ raise TypeError(f"Cannot assign object of type {type(value)} to DNDarray.")
+
+ original_key = key
+
+ self, processed_key = _resolve_indexing_state(
+ self, key, return_local_indices=True, op="set"
+ )
+
+ op = processed_key.op_type
+
+ # match dimensions (except for distr_mask as it perfectly aligns)
+ if op == "distr_mask":
+ value_is_scalar = (
+ np.isscalar(value)
+ or getattr(value, "ndim", 1) == 0
+ or (getattr(value, "shape", None) == (1,) and getattr(value, "split", 0) is None)
+ )
else:
- raise NotImplementedError(f"Not implemented for {value.__class__.__name__}")
+ value, value_is_scalar = self.__broadcast_value(
+ key, value, output_shape=processed_key.output_shape
+ )
+
+ # dispatch to the appropriate setter
+ if op == "distr_mask":
+ self.__setitem_mask(processed_key, original_key, value, value_is_scalar)
+ elif op == "scalar":
+ self.__setitem_scalar(processed_key, value, value_is_scalar)
+ elif op == "distributed":
+ self.__setitem_advanced_distributed(processed_key, original_key, value, value_is_scalar)
+ elif op == "slice":
+ self.__setitem_slice(processed_key, value, value_is_scalar)
+ elif op == "descending_slice":
+ self.__setitem_descending_slice_distributed(processed_key, value, value_is_scalar)
+ elif op in ("local_mask", "advanced"):
+ self.__setitem_advanced_local(processed_key, original_key, value, value_is_scalar)
def __str__(self) -> str:
"""
@@ -1890,7 +3324,7 @@ def __str__(self) -> str:
"""
return printing.__str__(self)
- def tolist(self, keepsplit: bool = False) -> List:
+ def tolist(self, keepsplit: bool = False) -> list:
"""
Return a copy of the local array data as a (nested) Python list. For scalars, a standard Python number is returned.
@@ -1936,39 +3370,20 @@ def __torch_function__(cls, func, types, args=(), kwargs=None):
def __torch_proxy__(self) -> torch.Tensor:
"""
- Return a 1-element `torch.Tensor` strided as the global `self` shape.
- Used internally for sanitation purposes.
+ Return a 1-element `torch.Tensor` strided as the global `self` shape. The split axis of the initial DNDarray is stored in the `names` attribute of the returned tensor.
+ Used internally to lower memory footprint of sanitation.
"""
- return torch.ones((1,), dtype=torch.int8, device=self.larray.device).as_strided(
- self.gshape, [0] * self.ndim
+ names = [None] * self.ndim
+ if self.split is not None:
+ names[self.split] = "split"
+ return (
+ torch.ones((1,), dtype=torch.int8, device=self.larray.device)
+ .as_strided(self.gshape, [0] * self.ndim)
+ .refine_names(*names)
)
- @staticmethod
- def __xitem_get_key_start_stop(
- rank: int,
- actives: list,
- key_st: int,
- key_sp: int,
- step: int,
- ends: torch.Tensor,
- og_key_st: int,
- ) -> Tuple[int, int]:
- # this does some basic logic for adjusting the starting and stoping of the a key for
- # setitem and getitem
- if step is not None and rank > actives[0]:
- offset = (ends[rank - 1] - og_key_st) % step
- if step > 2 and offset > 0:
- key_st += step - offset
- elif step == 2 and offset > 0:
- key_st += (ends[rank - 1] - og_key_st) % step
- if isinstance(key_st, torch.Tensor):
- key_st = key_st.item()
- if isinstance(key_sp, torch.Tensor):
- key_sp = key_sp.item()
- return key_st, key_sp
-
-
-# HeAT imports at the end to break cyclic dependencies
+
+# Heat imports at the end to break cyclic dependencies
from . import complex_math
from . import devices
from . import factories
@@ -1986,3 +3401,4 @@ def __xitem_get_key_start_stop(
from .devices import Device
from .stride_tricks import sanitize_axis
from .types import datatype, canonical_heat_type
+from .types import bool as ht_bool, uint8 as ht_uint8
diff --git a/heat/core/factories.py b/heat/core/factories.py
index a3c8f0a80c..6399bd9296 100644
--- a/heat/core/factories.py
+++ b/heat/core/factories.py
@@ -141,8 +141,9 @@ def arange(
else:
data = torch.arange(start, stop, step, device=device.torch_device)
data = data.type(htype.torch_type())
-
- return DNDarray(data, gshape, htype, split, device, comm, balanced)
+ return DNDarray(
+ data, gshape=gshape, dtype=htype, split=split, device=device, comm=comm, balanced=balanced
+ )
def array(
@@ -480,7 +481,15 @@ def array(
if gmatch != comm.size:
balanced = False
- return DNDarray(obj, tuple(gshape), dtype, split, device, comm, balanced)
+ return DNDarray(
+ obj,
+ gshape=tuple(gshape),
+ dtype=dtype,
+ split=split,
+ device=device,
+ comm=comm,
+ balanced=balanced,
+ )
def asarray(
@@ -725,7 +734,13 @@ def eye(
data = sanitize_memory_layout(data, order=order)
return DNDarray(
- data, gshape, types.canonical_heat_type(data.dtype), split, device, comm, balanced
+ data,
+ gshape=gshape,
+ dtype=types.canonical_heat_type(data.dtype),
+ split=split,
+ device=device,
+ comm=comm,
+ balanced=balanced,
)
@@ -780,7 +795,9 @@ def __factory(
data = local_factory(local_shape, dtype=dtype.torch_type(), device=device.torch_device)
data = sanitize_memory_layout(data, order=order)
- return DNDarray(data, shape, dtype, split, device, comm, balanced=True)
+ return DNDarray(
+ data, gshape=shape, dtype=dtype, split=split, device=device, comm=comm, balanced=True
+ )
def __factory_like(
@@ -1004,7 +1021,13 @@ def __from_partition_dict_helper(parted: dict, comm: Communication):
balanced = all(x[0][0] == x[1][0] for x in expected.values())
ret = DNDarray(
- data, gshape, htype, split, devices.sanitize_device(None), sanitize_comm(comm), balanced
+ data,
+ gshape=gshape,
+ dtype=htype,
+ split=split,
+ device=devices.sanitize_device(None),
+ comm=sanitize_comm(comm),
+ balanced=balanced,
)
ret.__partitions_dict__ = parted
@@ -1207,7 +1230,13 @@ def linspace(
# construct the resulting global tensor
ht_tensor = DNDarray(
- data, gshape, types.canonical_heat_type(data.dtype), split, device, comm, balanced
+ data,
+ gshape=gshape,
+ dtype=types.canonical_heat_type(data.dtype),
+ split=split,
+ device=device,
+ comm=comm,
+ balanced=balanced,
)
if retstep:
diff --git a/heat/core/indexing.py b/heat/core/indexing.py
index 916aa450df..96a6022eef 100644
--- a/heat/core/indexing.py
+++ b/heat/core/indexing.py
@@ -3,22 +3,23 @@
"""
import torch
-from typing import List, Dict, Any, TypeVar, Union, Tuple, Sequence
from .communication import MPI
from .dndarray import DNDarray
-from . import sanitation
+from . import factories
from . import types
+from . import manipulations
+from . import sanitation
__all__ = ["nonzero", "where"]
-def nonzero(x: DNDarray) -> DNDarray:
+def nonzero(x: DNDarray, as_tuple: bool = True) -> tuple[DNDarray, ...] | DNDarray:
"""
- Return a :class:`~heat.core.dndarray.DNDarray` containing the indices of the elements that are non-zero (using ``torch.nonzero``).
- If ``x`` is split then the result is split in the first dimension. However, this :class:`~heat.core.dndarray.DNDarray`
+ Return a Tuple of :class:`~heat.core.dndarray.DNDarray`s, one for each dimension of ``x``,
+ containing the indices of the non-zero elements in that dimension. If ``x`` is split then
+ the result is split in the 0th dimension. However, this :class:`~heat.core.dndarray.DNDarray`
can be UNBALANCED as it contains the indices of the non-zero elements on each node.
- Returns an array with one entry for each dimension of ``x``, containing the indices of the non-zero elements in that dimension.
The values in ``x`` are always tested and returned in row-major, C-style order.
The corresponding non-zero values can be obtained with: ``x[nonzero(x)]``.
@@ -26,16 +27,16 @@ def nonzero(x: DNDarray) -> DNDarray:
----------
x: DNDarray
Input array
+ as_tuple: bool, optional
+ Default is True for numpy-style nonzero output. If False, the output is a torch-style single 2D ``DNDarray`` of shape `(num_nonzero, ndim)` containing the indices of the non-zero elements.
Examples
--------
>>> import heat as ht
>>> x = ht.array([[3, 0, 0], [0, 4, 1], [0, 6, 0]], split=0)
>>> ht.nonzero(x)
- DNDarray([[0, 0],
- [1, 1],
- [1, 2],
- [2, 1]], dtype=ht.int64, device=cpu:0, split=0)
+ (DNDarray([0, 1, 1, 2], dtype=ht.int64, device=cpu:0, split=None),
+ DNDarray([0, 1, 2, 1], dtype=ht.int64, device=cpu:0, split=None))
>>> y = ht.array([[1, 2, 3], [4, 5, 6], [7, 8, 9]], split=0)
>>> y > 3
DNDarray([[False, False, False],
@@ -48,73 +49,123 @@ def nonzero(x: DNDarray) -> DNDarray:
[2, 0],
[2, 1],
[2, 2]], dtype=ht.int64, device=cpu:0, split=0)
+ (DNDarray([1, 1, 1, 2, 2, 2], dtype=ht.int64, device=cpu:0, split=None),
+ DNDarray([0, 1, 2, 0, 1, 2], dtype=ht.int64, device=cpu:0, split=None))
>>> y[ht.nonzero(y > 3)]
DNDarray([4, 5, 6, 7, 8, 9], dtype=ht.int64, device=cpu:0, split=0)
"""
sanitation.sanitize_in(x)
+ local_x = x.larray
+
+ if not x.is_distributed():
+ # nonzero indices as tuple
+ nonzero = torch.nonzero(input=local_x, as_tuple=as_tuple)
+ # ensure output split is consistent with distributed execution
+ out_split = 0 if x.split is not None else None
+
+ # bookkeeping for final DNDarray construct
+ if as_tuple:
+ nonzero = list(nonzero)
+ for i, nz_tensor in enumerate(nonzero):
+ nonzero[i] = factories.array(
+ nz_tensor, split=out_split, device=x.device, comm=x.comm
+ )
+ return tuple(nonzero)
+ # nonzero indices as single 2D DNDarray
+ return factories.array(nonzero, split=out_split, device=x.device, comm=x.comm)
+
+ # distributed case
+ lcl_nonzero = torch.nonzero(input=local_x, as_tuple=False)
+ nonzero_size = torch.tensor(lcl_nonzero.shape[0], dtype=torch.int64)
+ nonzero_dtype = types.canonical_heat_type(lcl_nonzero.dtype)
+
+ # global nonzero_size
+ x.comm.Allreduce(MPI.IN_PLACE, nonzero_size, MPI.SUM)
+ # correct indices along split axis
+ _, displs = x.counts_displs()
+ lcl_nonzero[:, x.split] += displs[x.comm.rank]
+
+ if x.split != 0:
+ # construct global 2D DNDarray of nz indices:
+ shape_2d = (nonzero_size.item(), x.ndim)
+ global_nonzero = DNDarray(
+ lcl_nonzero,
+ gshape=shape_2d,
+ dtype=nonzero_dtype,
+ split=0,
+ device=x.device,
+ comm=x.comm,
+ balanced=False,
+ )
+ # vectorized sorting of nz indices along axis 0
+ global_nonzero.balance_()
+ global_nonzero = manipulations.unique(global_nonzero, axis=0)
+ if not as_tuple:
+ # return indices as single 2D DNDarray
+ return global_nonzero
+ # return indices as tuple of 1D DNDarrays
+ lcl_nonzero = global_nonzero.larray.unbind(dim=1)
+ return tuple(
+ DNDarray(
+ nz_tensor,
+ gshape=(nonzero_size.item(),),
+ dtype=nonzero_dtype,
+ split=0,
+ device=x.device,
+ comm=x.comm,
+ balanced=True,
+ )
+ for nz_tensor in lcl_nonzero
+ )
- lcl_nonzero = torch.nonzero(input=x.larray, as_tuple=False)
-
- # add offsets mapping from local indices to global indices if x is split
- if x.split is not None:
- _, _, slices = x.comm.chunk(x.shape, x.split)
- lcl_nonzero[..., x.split] += slices[x.split].start
-
- if x.ndim == 1:
- lcl_nonzero = lcl_nonzero.squeeze(dim=1)
-
- # compute global shape of the index array
- gout = list(lcl_nonzero.shape)
- if x.split is None:
- is_split = None
- else:
- gout[0] = x.comm.allreduce(gout[0], MPI.SUM)
- is_split = 0
-
- return DNDarray(
- lcl_nonzero,
- gshape=tuple(gout),
- dtype=types.canonical_heat_type(lcl_nonzero.dtype),
- split=is_split,
- device=x.device,
- comm=x.comm,
- balanced=False,
+ # for split=0, the local nonzero indices are already globally ordered along the split axis
+ if not as_tuple:
+ # return indices as single 2D DNDarray
+ return DNDarray(
+ lcl_nonzero,
+ gshape=(nonzero_size.item(), x.ndim),
+ dtype=nonzero_dtype,
+ split=0,
+ device=x.device,
+ comm=x.comm,
+ balanced=False,
+ )
+ # return indices as tuple of 1D DNDarrays
+ lcl_nonzero = lcl_nonzero.unbind(dim=1)
+ return tuple(
+ DNDarray(
+ nz_tensor,
+ gshape=(nonzero_size.item(),),
+ dtype=nonzero_dtype,
+ split=0,
+ device=x.device,
+ comm=x.comm,
+ balanced=False,
+ )
+ for nz_tensor in lcl_nonzero
)
-DNDarray.nonzero = lambda self: nonzero(self)
+DNDarray.nonzero = lambda self: nonzero(self, as_tuple=True)
DNDarray.nonzero.__doc__ = nonzero.__doc__
def where(
- cond: DNDarray,
- x: Union[None, int, float, DNDarray] = None,
- y: Union[None, int, float, DNDarray] = None,
-) -> DNDarray:
+ cond: DNDarray, x: None | int | float | DNDarray = None, y: None | int | float | DNDarray = None
+) -> DNDarray | tuple[DNDarray, ...]:
"""
Return a :class:`~heat.core.dndarray.DNDarray` containing elements chosen from ``x`` or ``y`` depending on condition.
- Result is a :class:`~heat.core.dndarray.DNDarray` with elements from ``x`` where cond is ``True``,
- and elements from ``y`` elsewhere (``False``).
+ Result is a :class:`~heat.core.dndarray.DNDarray` with elements from ``x`` where ``cond`` is True, and from ``y`` elsewhere.
+
+ If only ``cond`` is provided, this function acts as a shorthand for :func:`nonzero`.
Parameters
----------
- cond : DNDarray
- Condition of interest, where true yield ``x`` otherwise yield ``y``
- x : DNDarray or int or float, optional
- Values from which to choose. ``x``, ``y`` and condition need to be broadcastable to some shape.
- y : DNDarray or int or float, optional
- Values from which to choose. ``x``, ``y`` and condition need to be broadcastable to some shape.
-
- Raises
- ------
- NotImplementedError
- if splits of the two input :class:`~heat.core.dndarray.DNDarray` differ
- TypeError
- if only x or y is given or both are not DNDarrays or numerical scalars
-
- Notes
- -----
- When only condition is provided, this function is a shorthand for :func:`nonzero`.
+ cond: DNDarray
+ When True, yield ``x``, otherwise yield ``y``.
+ x, y: DNDarray or scalar, optional
+ Values from which to choose. ``x``, ``y`` and ``cond`` must be broadcastable to some shape.
+ If ``x`` and ``y`` are distributed, they must have the same split axis as ``cond``.
Examples
--------
@@ -122,26 +173,39 @@ def where(
>>> x = ht.arange(10, split=0)
>>> ht.where(x < 5, x, 10 * x)
DNDarray([ 0, 1, 2, 3, 4, 50, 60, 70, 80, 90], dtype=ht.int64, device=cpu:0, split=0)
- >>> y = ht.array([[0, 1, 2], [0, 2, 4], [0, 3, 6]])
- >>> ht.where(y < 4, y, -1)
- DNDarray([[ 0, 1, 2],
- [ 0, 2, -1],
- [ 0, 3, -1]], dtype=ht.int64, device=cpu:0, split=None)
+
+ >>> # Indices retrieval (shorthand for nonzero)
+ >>> y = ht.array([[1, 2, 3], [4, 5, 6], [7, 8, 9]], split=0)
+ >>> ht.where(y > 3)
+ (DNDarray([1, 1, 1, 2, 2, 2], dtype=ht.int64, device=cpu:0, split=0),
+ DNDarray([0, 1, 2, 0, 1, 2], dtype=ht.int64, device=cpu:0, split=0))
"""
+ # ---- binary where(cond, x, y) branch ------------------------------------
if cond.split is not None and (isinstance(x, DNDarray) or isinstance(y, DNDarray)):
if (isinstance(x, DNDarray) and cond.split != x.split) or (
isinstance(y, DNDarray) and cond.split != y.split
):
- if len(y.shape) >= 1 and y.shape[0] > 1:
+ # Only raise if the "other" array has a meaningful first dimension.
+ if isinstance(y, DNDarray) and len(y.shape) >= 1 and y.shape[0] > 1:
raise NotImplementedError("binary op not implemented for different split axes")
+
if isinstance(x, (DNDarray, int, float)) and isinstance(y, (DNDarray, int, float)):
+ # Simple elementwise selection using arithmetic:
+ # cond == 0 -> take y, cond == 1 -> take x
for var in [x, y]:
if isinstance(var, int):
var = float(var)
return cond.dtype(cond == 0) * y + cond * x
+
+ # ---- where(cond) "indices only" branch ----------------------------------
elif x is None and y is None:
- return nonzero(cond)
+ # nonzero() properly handles all cases
+ nz = nonzero(cond)
+ return nz
+
+ # ---- invalid combinations ----------------------------------------------
else:
raise TypeError(
- f"either both or neither x and y must be given and both must be DNDarrays or numerical scalars({type(x)}, {type(y)})"
+ "either both or neither x and y must be given and both must be "
+ f"DNDarrays or numerical scalars (got {type(x)}, {type(y)})"
)
diff --git a/heat/core/linalg/basics.py b/heat/core/linalg/basics.py
index bd8240ad76..6614817cd3 100644
--- a/heat/core/linalg/basics.py
+++ b/heat/core/linalg/basics.py
@@ -2347,12 +2347,12 @@ def transpose(a: DNDarray, axes: Optional[List[int]] = None) -> DNDarray:
return DNDarray(
transposed_data,
- transposed_shape,
- a.dtype,
- transposed_split,
- a.device,
- a.comm,
- a.balanced,
+ gshape=transposed_shape,
+ dtype=a.dtype,
+ split=transposed_split,
+ device=a.device,
+ comm=a.comm,
+ balanced=a.balanced,
)
# if not possible re- raise any torch exception as ValueError
except (RuntimeError, IndexError) as exception:
diff --git a/heat/core/manipulations.py b/heat/core/manipulations.py
index 200120fdad..29b8e442b2 100644
--- a/heat/core/manipulations.py
+++ b/heat/core/manipulations.py
@@ -4077,7 +4077,7 @@ def tile(x: DNDarray, reps: Sequence[int, ...]) -> DNDarray:
except AttributeError:
x = factories.array(x).reshape(1)
- x_proxy = x.__torch_proxy__()
+ x_proxy = x.__torch_proxy__().rename(None) # drop named-tensor metadata
# torch-proof args/kwargs:
# torch `reps`: int or sequence of ints; numpy `reps`: can be array-like
@@ -4141,7 +4141,7 @@ def tile(x: DNDarray, reps: Sequence[int, ...]) -> DNDarray:
trans_axes[0], trans_axes[x.split] = x.split, 0
reps[0], reps[x.split] = reps[x.split], reps[0]
x = linalg.transpose(x, trans_axes)
- x_proxy = x.__torch_proxy__()
+ x_proxy = x.__torch_proxy__().rename(None)
out_gshape = tuple(x_proxy.repeat(reps).shape)
local_x = x.larray
diff --git a/heat/core/memory.py b/heat/core/memory.py
index dbf2d8723e..b18ab9e5ec 100644
--- a/heat/core/memory.py
+++ b/heat/core/memory.py
@@ -32,7 +32,15 @@ def copy(x: DNDarray) -> DNDarray:
DNDarray([1, 2, 3], dtype=ht.int64, device=cpu:0, split=None)
"""
sanitation.sanitize_in(x)
- return DNDarray(x.larray.clone(), x.shape, x.dtype, x.split, x.device, x.comm, x.balanced)
+ return DNDarray(
+ x.larray.clone(),
+ gshape=x.gshape,
+ dtype=x.dtype,
+ split=x.split,
+ device=x.device,
+ comm=x.comm,
+ balanced=x.balanced,
+ )
DNDarray.copy = lambda self: copy(self)
diff --git a/tests/cluster/test_kmedians.py b/tests/cluster/test_kmedians.py
index dda91562ee..573ac47e7e 100644
--- a/tests/cluster/test_kmedians.py
+++ b/tests/cluster/test_kmedians.py
@@ -41,7 +41,7 @@ def test_fit_iris_unsplit(self):
# fit the clusters
k = 3
- kmedian = ht.cluster.KMedians(n_clusters=k)
+ kmedian = ht.cluster.KMedians(n_clusters=k, random_state=1)
kmedian.fit(iris)
# check whether the results are correct
diff --git a/tests/core/linalg/test_basics.py b/tests/core/linalg/test_basics.py
index 21c0ce614e..4ec900f555 100644
--- a/tests/core/linalg/test_basics.py
+++ b/tests/core/linalg/test_basics.py
@@ -301,6 +301,8 @@ def test_inv(self):
self.assertTupleEqual(ainv.shape, a.shape)
self.assertTrue(ht.allclose(ainv, ares, atol=1e-6))
+ # distributed
+ # ares = ht.array([[2.0, 2, 1], [3, 4, 1], [0, 1, -1]], split=0)
a = ht.array([[5.0, -3, 2], [-3, 2, -1], [-3, 2, -2]], split=0)
ainv = ht.linalg.inv(a)
self.assertEqual(ainv.split, a.split)
@@ -401,12 +403,8 @@ def test_matmul(self):
b_torch[:, 0] = torch.arange(1, j + 1, device=self.device.torch_device)
# splits None None
- a = ht.ones((n, m), split=None)
- b = ht.ones((j, k), split=None)
- a[0] = ht.arange(1, m + 1)
- a[:, -1] = ht.arange(1, n + 1)
- b[0] = ht.arange(1, k + 1)
- b[:, 0] = ht.arange(1, j + 1)
+ a = ht.array(a_torch, split=None, copy=True)
+ b = ht.array(b_torch, split=None, copy=True)
ret00 = ht.matmul(a, b)
self.assertEqual(ht.all(ret00 == ht.array(a_torch @ b_torch)), 1)
@@ -418,12 +416,8 @@ def test_matmul(self):
self.assertEqual(b.split, None)
# splits None None
- a = ht.ones((n, m), split=None)
- b = ht.ones((j, k), split=None)
- a[0] = ht.arange(1, m + 1)
- a[:, -1] = ht.arange(1, n + 1)
- b[0] = ht.arange(1, k + 1)
- b[:, 0] = ht.arange(1, j + 1)
+ a = ht.array(a_torch, split=None, copy=True)
+ b = ht.array(b_torch, split=None, copy=True)
ret00 = ht.matmul(a, b, allow_resplit=True)
self.assertEqual(ht.all(ret00 == ht.array(a_torch @ b_torch)), 1)
@@ -437,12 +431,8 @@ def test_matmul(self):
# splits 0 None on 1 process
if a.comm.size == 1:
- a = ht.ones((n, m), split=0)
- b = ht.ones((j, k), split=None)
- a[0] = ht.arange(1, m + 1)
- a[:, -1] = ht.arange(1, n + 1)
- b[0] = ht.arange(1, k + 1)
- b[:, 0] = ht.arange(1, j + 1)
+ a = ht.array(a_torch, split=0, copy=True)
+ b = ht.array(b_torch, split=None, copy=True)
ret00 = ht.matmul(a, b, allow_resplit=True)
self.assertEqual(ht.all(ret00 == ht.array(a_torch @ b_torch)), 1)
@@ -455,12 +445,8 @@ def test_matmul(self):
if a.comm.size > 1:
# splits 00
- a = ht.ones((n, m), split=0, dtype=ht.float64)
- b = ht.ones((j, k), split=0)
- a[0] = ht.arange(1, m + 1)
- a[:, -1] = ht.arange(1, n + 1)
- b[0] = ht.arange(1, k + 1)
- b[:, 0] = ht.arange(1, j + 1)
+ a = ht.array(a_torch, split=0, dtype=ht.float64, copy=True)
+ b = ht.array(b_torch, split=0, copy=True)
ret00 = a @ b
ret_comp00 = ht.array(a_torch @ b_torch, split=0)
@@ -487,12 +473,8 @@ def test_matmul(self):
self.assertEqual(ret00.split, 0)
# splits 01
- a = ht.ones((n, m), split=0)
- b = ht.ones((j, k), split=1, dtype=ht.float64)
- a[0] = ht.arange(1, m + 1)
- a[:, -1] = ht.arange(1, n + 1)
- b[0] = ht.arange(1, k + 1)
- b[:, 0] = ht.arange(1, j + 1)
+ a = ht.array(a_torch, split=0, copy=True)
+ b = ht.array(b_torch, split=1, dtype=ht.float64, copy=True)
ret00 = ht.matmul(a, b)
ret_comp01 = ht.array(a_torch @ b_torch, split=0)
@@ -503,12 +485,8 @@ def test_matmul(self):
self.assertEqual(ret00.split, 0)
# splits 10
- a = ht.ones((n, m), split=1)
- b = ht.ones((j, k), split=0)
- a[0] = ht.arange(1, m + 1)
- a[:, -1] = ht.arange(1, n + 1)
- b[0] = ht.arange(1, k + 1)
- b[:, 0] = ht.arange(1, j + 1)
+ a = ht.array(a_torch, split=1, copy=True)
+ b = ht.array(b_torch, split=0, copy=True)
ret00 = ht.matmul(a, b)
ret_comp10 = ht.array(a_torch @ b_torch, split=1)
@@ -519,28 +497,8 @@ def test_matmul(self):
self.assertEqual(ret00.split, 1)
# splits 11
- a = ht.ones((n, m), split=1)
- b = ht.ones((j, k), split=1)
- a[0] = ht.arange(1, m + 1)
- a[:, -1] = ht.arange(1, n + 1)
- b[0] = ht.arange(1, k + 1)
- b[:, 0] = ht.arange(1, j + 1)
- ret00 = ht.matmul(a, b)
-
- ret_comp11 = ht.array(a_torch @ b_torch, split=1)
- self.assertTrue(ht.equal(ret00, ret_comp11))
- self.assertIsInstance(ret00, ht.DNDarray)
- self.assertEqual(ret00.shape, (n, k))
- self.assertEqual(ret00.dtype, ht.float)
- self.assertEqual(ret00.split, 1)
-
- # splits 11 (torch)
- a = ht.array(torch.ones((n, m), device=self.device.torch_device), split=1)
- b = ht.array(torch.ones((j, k), device=self.device.torch_device), split=1)
- a[0] = ht.arange(1, m + 1)
- a[:, -1] = ht.arange(1, n + 1)
- b[0] = ht.arange(1, k + 1)
- b[:, 0] = ht.arange(1, j + 1)
+ a = ht.array(a_torch, split=1, copy=True)
+ b = ht.array(b_torch, split=1, copy=True)
ret00 = ht.matmul(a, b)
ret_comp11 = ht.array(a_torch @ b_torch, split=1)
@@ -551,12 +509,8 @@ def test_matmul(self):
self.assertEqual(ret00.split, 1)
# splits 0 None
- a = ht.ones((n, m), split=0)
- b = ht.ones((j, k), split=None)
- a[0] = ht.arange(1, m + 1)
- a[:, -1] = ht.arange(1, n + 1)
- b[0] = ht.arange(1, k + 1)
- b[:, 0] = ht.arange(1, j + 1)
+ a = ht.array(a_torch, split=0, copy=True)
+ b = ht.array(b_torch, split=None, copy=True)
ret00 = ht.matmul(a, b)
ret_comp0 = ht.array(a_torch @ b_torch, split=0)
@@ -566,13 +520,10 @@ def test_matmul(self):
self.assertEqual(ret00.dtype, ht.float)
self.assertEqual(ret00.split, 0)
+
# splits 1 None
- a = ht.ones((n, m), split=1)
- b = ht.ones((j, k), split=None)
- a[0] = ht.arange(1, m + 1)
- a[:, -1] = ht.arange(1, n + 1)
- b[0] = ht.arange(1, k + 1)
- b[:, 0] = ht.arange(1, j + 1)
+ a = ht.array(a_torch, split=1, copy=True)
+ b = ht.array(b_torch, split=None, copy=True)
ret00 = ht.matmul(a, b)
ret_comp1 = ht.array(a_torch @ b_torch, split=1)
@@ -583,12 +534,8 @@ def test_matmul(self):
self.assertEqual(ret00.split, 1)
# splits None 0
- a = ht.ones((n, m), split=None)
- b = ht.ones((j, k), split=0)
- a[0] = ht.arange(1, m + 1)
- a[:, -1] = ht.arange(1, n + 1)
- b[0] = ht.arange(1, k + 1)
- b[:, 0] = ht.arange(1, j + 1)
+ a = ht.array(a_torch, split=None, copy=True)
+ b = ht.array(b_torch, split=0, copy=True)
ret00 = ht.matmul(a, b)
ret_comp = ht.array(a_torch @ b_torch, split=0)
@@ -599,12 +546,8 @@ def test_matmul(self):
self.assertEqual(ret00.split, 0)
# splits None 1
- a = ht.ones((n, m), split=None)
- b = ht.ones((j, k), split=1)
- a[0] = ht.arange(1, m + 1)
- a[:, -1] = ht.arange(1, n + 1)
- b[0] = ht.arange(1, k + 1)
- b[:, 0] = ht.arange(1, j + 1)
+ a = ht.array(a_torch, split=None, copy=True)
+ b = ht.array(b_torch, split=1, copy=True)
ret00 = ht.matmul(a, b)
ret_comp = ht.array(a_torch @ b_torch, split=1)
@@ -621,10 +564,8 @@ def test_matmul(self):
b_torch[0] = torch.arange(1, k + 1, device=self.device.torch_device)
b_torch[:, 0] = torch.arange(1, j + 1, device=self.device.torch_device)
# splits None None
- a = ht.ones((m), split=None)
- b = ht.ones((j, k), split=None)
- b[0] = ht.arange(1, k + 1)
- b[:, 0] = ht.arange(1, j + 1)
+ a = ht.array(a_torch, split=None, copy=True)
+ b = ht.array(b_torch, split=None, copy=True)
ret00 = ht.matmul(a, b)
ret_comp = ht.array(a_torch @ b_torch, split=None)
@@ -636,10 +577,8 @@ def test_matmul(self):
self.assertEqual(ret00.split, None)
# splits None 0
- a = ht.ones((m), split=None)
- b = ht.ones((j, k), split=0)
- b[0] = ht.arange(1, k + 1)
- b[:, 0] = ht.arange(1, j + 1)
+ a = ht.array(a_torch, split=None, copy=True)
+ b = ht.array(b_torch, split=0, copy=True)
ret00 = ht.matmul(a, b)
ret_comp = ht.array(a_torch @ b_torch, split=None)
@@ -650,10 +589,8 @@ def test_matmul(self):
self.assertEqual(ret00.split, 0)
# splits None 1
- a = ht.ones((m), split=None)
- b = ht.ones((j, k), split=1)
- b[0] = ht.arange(1, k + 1)
- b[:, 0] = ht.arange(1, j + 1)
+ a = ht.array(a_torch, split=None, copy=True)
+ b = ht.array(b_torch, split=1, copy=True)
ret00 = ht.matmul(a, b)
ret_comp = ht.array(a_torch @ b_torch, split=0)
self.assertTrue(ht.equal(ret00, ret_comp))
@@ -663,10 +600,8 @@ def test_matmul(self):
self.assertEqual(ret00.split, 0)
# splits 0 None
- a = ht.ones((m), split=None)
- b = ht.ones((j, k), split=0)
- b[0] = ht.arange(1, k + 1)
- b[:, 0] = ht.arange(1, j + 1)
+ a = ht.array(a_torch, split=None, copy=True)
+ b = ht.array(b_torch, split=0, copy=True)
ret00 = ht.matmul(a, b)
ret_comp = ht.array(a_torch @ b_torch, split=None)
@@ -677,10 +612,8 @@ def test_matmul(self):
self.assertEqual(ret00.split, 0)
# splits 0 0
- a = ht.ones((m), split=0)
- b = ht.ones((j, k), split=0)
- b[0] = ht.arange(1, k + 1)
- b[:, 0] = ht.arange(1, j + 1)
+ a = ht.array(a_torch, split=0, copy=True)
+ b = ht.array(b_torch, split=0, copy=True)
ret00 = ht.matmul(a, b)
ret_comp = ht.array(a_torch @ b_torch, split=None)
@@ -691,10 +624,8 @@ def test_matmul(self):
self.assertEqual(ret00.split, 0)
# splits 0 1
- a = ht.ones((m), split=0)
- b = ht.ones((j, k), split=1)
- b[0] = ht.arange(1, k + 1)
- b[:, 0] = ht.arange(1, j + 1)
+ a = ht.array(a_torch, split=0, copy=True)
+ b = ht.array(b_torch, split=1, copy=True)
ret00 = ht.matmul(a, b)
ret_comp = ht.array(a_torch @ b_torch, split=None)
@@ -710,10 +641,8 @@ def test_matmul(self):
a_torch[:, -1] = torch.arange(1, n + 1, device=self.device.torch_device)
b_torch = torch.ones((j), device=self.device.torch_device)
# splits None None
- a = ht.ones((n, m), split=None)
- b = ht.ones((j), split=None)
- a[0] = ht.arange(1, m + 1)
- a[:, -1] = ht.arange(1, n + 1)
+ a = ht.array(a_torch, split=None, copy=True)
+ b = ht.array(b_torch, split=None, copy=True)
ret00 = ht.matmul(a, b)
ret_comp = ht.array(a_torch @ b_torch, split=None)
@@ -723,10 +652,8 @@ def test_matmul(self):
self.assertEqual(ret00.dtype, ht.float)
self.assertEqual(ret00.split, None)
- a = ht.ones((n, m), split=None, dtype=ht.int64)
- b = ht.ones((j), split=None, dtype=ht.int64)
- a[0] = ht.arange(1, m + 1, dtype=ht.int64)
- a[:, -1] = ht.arange(1, n + 1, dtype=ht.int64)
+ a = ht.array(a_torch, split=None, dtype=ht.int64, copy=True)
+ b = ht.array(b_torch, split=None, dtype=ht.int64, copy=True)
ret00 = ht.matmul(a, b)
ret_comp = ht.array((a_torch @ b_torch), split=None)
@@ -737,10 +664,8 @@ def test_matmul(self):
self.assertEqual(ret00.split, None)
# splits 0 None
- a = ht.ones((n, m), split=0)
- b = ht.ones((j), split=None)
- a[0] = ht.arange(1, m + 1)
- a[:, -1] = ht.arange(1, n + 1)
+ a = ht.array(a_torch, split=0, copy=True)
+ b = ht.array(b_torch, split=None, copy=True)
ret00 = ht.matmul(a, b)
ret_comp = ht.array((a_torch @ b_torch), split=None)
@@ -750,10 +675,8 @@ def test_matmul(self):
self.assertEqual(ret00.dtype, ht.float)
self.assertEqual(ret00.split, 0)
- a = ht.ones((n, m), split=0, dtype=ht.int64)
- b = ht.ones((j), split=None, dtype=ht.int64)
- a[0] = ht.arange(1, m + 1, dtype=ht.int64)
- a[:, -1] = ht.arange(1, n + 1, dtype=ht.int64)
+ a = ht.array(a_torch, split=0, dtype=ht.int64, copy=True)
+ b = ht.array(b_torch, split=None, dtype=ht.int64, copy=True)
ret00 = ht.matmul(a, b)
ret_comp = ht.array((a_torch @ b_torch), split=None)
@@ -764,10 +687,8 @@ def test_matmul(self):
self.assertEqual(ret00.split, 0)
# splits 1 None
- a = ht.ones((n, m), split=1)
- b = ht.ones((j), split=None)
- a[0] = ht.arange(1, m + 1)
- a[:, -1] = ht.arange(1, n + 1)
+ a = ht.array(a_torch, split=1, copy=True)
+ b = ht.array(b_torch, split=None, copy=True)
ret00 = ht.matmul(a, b)
ret_comp = ht.array((a_torch @ b_torch), split=None)
@@ -777,10 +698,8 @@ def test_matmul(self):
self.assertEqual(ret00.dtype, ht.float)
self.assertEqual(ret00.split, 0)
- a = ht.ones((n, m), split=1, dtype=ht.int64)
- b = ht.ones((j), split=None, dtype=ht.int64)
- a[0] = ht.arange(1, m + 1, dtype=ht.int64)
- a[:, -1] = ht.arange(1, n + 1, dtype=ht.int64)
+ a = ht.array(a_torch, split=1, dtype=ht.int64, copy=True)
+ b = ht.array(b_torch, split=None, dtype=ht.int64, copy=True)
ret00 = ht.matmul(a, b)
ret_comp = ht.array((a_torch @ b_torch), split=None)
@@ -791,10 +710,8 @@ def test_matmul(self):
self.assertEqual(ret00.split, 0)
# splits None 0
- a = ht.ones((n, m), split=None)
- b = ht.ones((j), split=0)
- a[0] = ht.arange(1, m + 1)
- a[:, -1] = ht.arange(1, n + 1)
+ a = ht.array(a_torch, split=None, copy=True)
+ b = ht.array(b_torch, split=0, copy=True)
ret00 = ht.matmul(a, b)
ret_comp = ht.array((a_torch @ b_torch), split=None)
@@ -804,10 +721,8 @@ def test_matmul(self):
self.assertEqual(ret00.dtype, ht.float)
self.assertEqual(ret00.split, 0)
- a = ht.ones((n, m), split=None, dtype=ht.int64)
- b = ht.ones((j), split=0, dtype=ht.int64)
- a[0] = ht.arange(1, m + 1, dtype=ht.int64)
- a[:, -1] = ht.arange(1, n + 1, dtype=ht.int64)
+ a = ht.array(a_torch, split=None, dtype=ht.int64, copy=True)
+ b = ht.array(b_torch, split=0, dtype=ht.int64, copy=True)
ret00 = ht.matmul(a, b)
ret_comp = ht.array((a_torch @ b_torch), split=None)
@@ -818,10 +733,8 @@ def test_matmul(self):
self.assertEqual(ret00.split, 0)
# splits 0 0
- a = ht.ones((n, m), split=0)
- b = ht.ones((j), split=0)
- a[0] = ht.arange(1, m + 1)
- a[:, -1] = ht.arange(1, n + 1)
+ a = ht.array(a_torch, split=0, copy=True)
+ b = ht.array(b_torch, split=0, copy=True)
ret00 = ht.matmul(a, b)
ret_comp = ht.array((a_torch @ b_torch), split=None)
@@ -831,10 +744,8 @@ def test_matmul(self):
self.assertEqual(ret00.dtype, ht.float)
self.assertEqual(ret00.split, 0)
- a = ht.ones((n, m), split=0, dtype=ht.int64)
- b = ht.ones((j), split=0, dtype=ht.int64)
- a[0] = ht.arange(1, m + 1, dtype=ht.int64)
- a[:, -1] = ht.arange(1, n + 1, dtype=ht.int64)
+ a = ht.array(a_torch, split=0, dtype=ht.int64, copy=True)
+ b = ht.array(b_torch, split=0, dtype=ht.int64, copy=True)
ret00 = ht.matmul(a, b)
ret_comp = ht.array((a_torch @ b_torch), split=None)
@@ -845,10 +756,8 @@ def test_matmul(self):
self.assertEqual(ret00.split, 0)
# splits 1 0
- a = ht.ones((n, m), split=1)
- b = ht.ones((j), split=0)
- a[0] = ht.arange(1, m + 1)
- a[:, -1] = ht.arange(1, n + 1)
+ a = ht.array(a_torch, split=1, copy=True)
+ b = ht.array(b_torch, split=0, copy=True)
ret00 = ht.matmul(a, b)
ret_comp = ht.array((a_torch @ b_torch), split=None)
@@ -858,10 +767,8 @@ def test_matmul(self):
self.assertEqual(ret00.dtype, ht.float)
self.assertEqual(ret00.split, 0)
- a = ht.ones((n, m), split=1, dtype=ht.int64)
- b = ht.ones((j), split=0, dtype=ht.int64)
- a[0] = ht.arange(1, m + 1, dtype=ht.int64)
- a[:, -1] = ht.arange(1, n + 1, dtype=ht.int64)
+ a = ht.array(a_torch, split=1, dtype=ht.int64, copy=True)
+ b = ht.array(b_torch, split=0, dtype=ht.int64, copy=True)
ret00 = ht.matmul(a, b)
ret_comp = ht.array((a_torch @ b_torch), split=None)
diff --git a/tests/core/test_dndarray.py b/tests/core/test_dndarray.py
index 7ca42be55d..2865c56d60 100644
--- a/tests/core/test_dndarray.py
+++ b/tests/core/test_dndarray.py
@@ -16,9 +16,9 @@ def setUpClass(cls):
N = ht.MPI_WORLD.size
cls.reference_tensor = ht.zeros((N, N + 1, 2 * N))
- for n in range(N):
- for m in range(N + 1):
- cls.reference_tensor[n, m, :] = ht.arange(0, 2 * N) + m * 10 + n * 100
+ # for n in range(N):
+ # for m in range(N + 1):
+ # cls.reference_tensor[n, m, :] = ht.arange(0, 2 * N) + m * 10 + n * 100
def test_and(self):
int16_tensor = ht.array([[1, 1], [2, 2]], dtype=ht.int16)
@@ -707,6 +707,341 @@ def test_float_cast(self):
with self.assertRaises(TypeError):
float(ht.full((ht.MPI_WORLD.size,), 2, split=0))
+ # indexing tests inspired by https://numpy.org/doc/stable/user/basics.indexing.html
+ def test_getitem_single_element(self):
+ # Single element indexing
+ # 1D, local
+ x = ht.arange(10)
+ self.assertTrue(x[2].item() == 2)
+ self.assertTrue(x[-2].item() == 8)
+ self.assertTrue(x[2].dtype == ht.int32)
+ # 1D, distributed
+ x = ht.arange(10, split=0, dtype=ht.float64)
+ self.assertTrue(x[2].item() == 2.0)
+ self.assertTrue(x[-2].item() == 8.0)
+ self.assertTrue(x[2].dtype == ht.float64)
+ self.assertTrue(x[2].split is None)
+ # 2D, local
+ x = ht.arange(10).reshape(2, 5)
+ self.assertTrue((x[0] == ht.arange(5)).all().item())
+ self.assertTrue(x[0].dtype == ht.int32)
+ # 2D, distributed
+ x_split0 = ht.array(x, split=0)
+ self.assertTrue((x_split0[0] == ht.arange(5, split=None)).all().item())
+ x_split1 = ht.array(x, split=1)
+ self.assertTrue((x_split1[-2] == ht.arange(5, split=0)).all().item())
+ # 3D, local
+ x = ht.arange(27).reshape(3, 3, 3)
+ key = -2
+ indexed = x[key]
+ self.assertTrue((indexed.larray == x.larray[key]).all())
+ self.assertTrue(indexed.dtype == ht.int32)
+ self.assertTrue(indexed.split is None)
+ # 3D, distributed, split = 0
+ x_split0 = ht.array(x, dtype=ht.float32, split=0)
+ indexed_split0 = x_split0[key]
+ self.assertTrue((indexed_split0.larray == x.larray[key]).all())
+ self.assertTrue(indexed_split0.dtype == ht.float32)
+ self.assertTrue(indexed_split0.split is None)
+ # 3D, distributed split, != 0
+ x_split2 = ht.array(x, dtype=ht.int64, split=2)
+ key = ht.array(2)
+ indexed_split2 = x_split2[key]
+ self.assertTrue((indexed_split2.numpy() == x.numpy()[key.item()]).all())
+ self.assertTrue(indexed_split2.dtype == ht.int64)
+ self.assertTrue(indexed_split2.split == 1)
+
+ # tests for bug 730:
+ a = ht.ones((10, 25, 30), split=1)
+ if a.comm.size > 1:
+ self.assertEqual(a[0].split, 0)
+ self.assertEqual(a[:, 0, :].split, None)
+ self.assertEqual(a[:, :, 0].split, 1)
+
+ def test_getitem_slicing(self):
+ # Slicing and striding
+ x = ht.arange(20, split=0)
+ x_sliced = x[1:11:3]
+ x_np = np.arange(20)
+ x_sliced_np = x_np[1:11:3]
+ self.assert_array_equal(x_sliced, x_sliced_np)
+ self.assertTrue(x_sliced.split == 0)
+
+ # 1-element slice along split axis
+ x = ht.arange(20).reshape(4, 5)
+ x.resplit_(axis=1)
+ x_sliced = x[:, 2:3]
+ x_np = np.arange(20).reshape(4, 5)
+ x_sliced_np = x_np[:, 2:3]
+ self.assert_array_equal(x_sliced, x_sliced_np)
+ self.assertTrue(x_sliced.split == 1)
+
+ def test_getitem_slicing_negative_step(self):
+ # slicing with negative step along split axis 0
+ shape = (20, 4, 3)
+ x_3d = ht.arange(20 * 4 * 3, split=0).reshape(shape)
+ x_3d_sliced = x_3d[17:2:-2, :2, ht.array(1)]
+ x_3d_sliced_np = np.arange(20 * 4 * 3).reshape(shape)[17:2:-2, :2, 1]
+ self.assert_array_equal(x_3d_sliced, x_3d_sliced_np)
+ self.assertTrue(x_3d_sliced.split == 0)
+
+ # slicing with negative step along split 1
+ shape = (4, 20, 3)
+ x_3d = ht.arange(20 * 4 * 3).reshape(shape)
+ x_3d.resplit_(axis=1)
+ key = (slice(None, 2), slice(17, 2, -2), 1)
+ x_3d_sliced = x_3d[key]
+ x_3d_sliced_np = np.arange(20 * 4 * 3).reshape(shape)[:2, 17:2:-2, 1]
+ self.assert_array_equal(x_3d_sliced, x_3d_sliced_np)
+ self.assertTrue(x_3d_sliced.split == 1)
+
+ # slicing with negative step along split 2 and loss of axis < split
+ shape = (4, 3, 20)
+ x_3d = ht.arange(20 * 4 * 3).reshape(shape)
+ x_3d.resplit_(axis=2)
+ key = (slice(None, 2), 1, slice(17, 10, -2))
+ x_3d_sliced = x_3d[key]
+ x_3d_sliced_np = np.arange(20 * 4 * 3).reshape(shape)[:2, 1, 17:10:-2]
+ self.assert_array_equal(x_3d_sliced, x_3d_sliced_np)
+ self.assertTrue(x_3d_sliced.split == 1)
+
+ # slicing with negative step along split 2 and loss of all axes but split
+ shape = (4, 3, 20)
+ x_3d = ht.arange(20 * 4 * 3).reshape(shape)
+ x_3d.resplit_(axis=2)
+ key = (0, 1, slice(17, 13, -1))
+ x_3d_sliced = x_3d[key]
+ x_3d_sliced_np = np.arange(20 * 4 * 3).reshape(shape)[0, 1, 17:13:-1]
+ self.assert_array_equal(x_3d_sliced, x_3d_sliced_np)
+ self.assertTrue(x_3d_sliced.split == 0)
+
+ def test_getitem_dimensional_indexing(self):
+ # ellipsis
+ x_np = np.array([[[1], [2], [3]], [[4], [5], [6]]])
+ x_np_ellipsis = x_np[..., 0]
+ x = ht.array([[[1], [2], [3]], [[4], [5], [6]]])
+
+ # local
+ x_ellipsis = x[..., 0]
+ x_slice = x[:, :, 0]
+ self.assert_array_equal(x_ellipsis, x_np_ellipsis)
+ self.assert_array_equal(x_slice, x_np_ellipsis)
+
+ # distributed
+ x.resplit_(axis=1)
+ x_ellipsis = x[..., 0]
+ x_slice = x[:, :, 0]
+ self.assert_array_equal(x_ellipsis, x_np_ellipsis)
+ self.assert_array_equal(x_slice, x_np_ellipsis)
+ self.assertTrue(x_ellipsis.split == 1)
+
+ # newaxis: local
+ x = ht.array([[[1], [2], [3]], [[4], [5], [6]]])
+ x_np_newaxis = x_np[:, np.newaxis, :2, :]
+ x_newaxis = x[:, np.newaxis, :2, :]
+ x_none = x[:, None, :2, :]
+ self.assert_array_equal(x_newaxis, x_np_newaxis)
+ self.assert_array_equal(x_none, x_np_newaxis)
+
+ # newaxis: distributed
+ x.resplit_(axis=1)
+ x_newaxis = x[:, np.newaxis, :2, :]
+ x_none = x[:, None, :2, :]
+ self.assert_array_equal(x_newaxis, x_np_newaxis)
+ self.assert_array_equal(x_none, x_np_newaxis)
+ self.assertTrue(x_newaxis.split == 2)
+ self.assertTrue(x_none.split == 2)
+
+ x = ht.arange(5, split=0)
+ x_np = np.arange(5)
+ y = x[:, np.newaxis] + x[np.newaxis, :]
+ y_np = x_np[:, np.newaxis] + x_np[np.newaxis, :]
+ self.assert_array_equal(y, y_np)
+ self.assertTrue(y.split == 0)
+
+ for split in [None, 0, 1, 2]:
+ for new_dim in [0, 1, 2]:
+ for add in [np.newaxis, None]:
+ arr = ht.ones((4, 3, 2), split=split, dtype=ht.int32)
+ check = torch.ones((4, 3, 2), dtype=torch.int32)
+ idx = [slice(None), slice(None), slice(None)]
+ idx[new_dim] = add
+ idx = tuple(idx)
+ arr = arr[idx]
+ check = check[idx]
+ self.assertTrue(arr.shape == check.shape)
+ self.assertTrue(arr.lshape[new_dim] == 1)
+
+ # test multiple ellipses rejection
+ a = ht.ones((5, 5))
+ with self.assertRaises(ValueError):
+ a[..., ...]
+
+ def test_getitem_advanced_indexing(self):
+ # "x[(1, 2, 3),] is fundamentally different from x[(1, 2, 3)]" cf. numpy docs
+
+ x_np = np.arange(60).reshape(5, 3, 4)
+ indexed_x_np = x_np[(1, 2, 3)]
+ adv_indexed_x_np = x_np[(1, 2, 3),]
+ x = ht.array(x_np, split=0)
+ indexed_x = x[(1, 2, 3)]
+ self.assertTrue(indexed_x.item() == np.array(indexed_x_np))
+ adv_indexed_x = x[(1, 2, 3),]
+ self.assert_array_equal(adv_indexed_x, adv_indexed_x_np)
+
+ # 1d
+ x = ht.arange(10, 1, -1, split=0)
+ x_np = np.arange(10, 1, -1)
+ x_adv_ind = x[np.array([3, 3, 1, 8])]
+ x_np_adv_ind = x_np[np.array([3, 3, 1, 8])]
+ self.assert_array_equal(x_adv_ind, x_np_adv_ind)
+
+ # 1d, split 0, advanced indexing with a local DNDarray
+ x = ht.arange(10, 1, -1, split=0)
+ x_np = np.arange(10, 1, -1)
+ idx_np = np.array([3, 3, 1, 8])
+ # local DNDarray index
+ idx = ht.array(idx_np, split=None)
+ x_adv_ind = x[idx]
+ x_np_adv_ind = x_np[idx_np]
+ self.assert_array_equal(x_adv_ind, x_np_adv_ind)
+
+ # 3d, split 0, non-unique, non-ordered key along split axis
+ x = ht.arange(60, split=0).reshape(5, 3, 4)
+ x_np = np.arange(60).reshape(5, 3, 4)
+ k1 = np.array([0, 4, 1, 0])
+ k2 = np.array([0, 2, 1, 0])
+ k3 = np.array([1, 2, 3, 1])
+ self.assert_array_equal(
+ x[ht.array(k1, split=0), ht.array(k2, split=0), ht.array(k3, split=0)], x_np[k1, k2, k3]
+ )
+ # advanced indexing on non-consecutive dimensions
+ x = ht.arange(60, split=0).reshape(5, 3, 4, new_split=1)
+ x_copy = x.copy()
+ x_np = np.arange(60).reshape(5, 3, 4)
+ k1 = np.array([0, 4, 1, 0])
+ k2 = 0
+ k3 = np.array([1, 2, 3, 1])
+ key = (k1, k2, k3)
+ self.assert_array_equal(x[key], x_np[key])
+ # check that x is unchanged after internal manipulation
+ self.assertTrue(x.shape == x_copy.shape)
+ self.assertTrue(x.split == x_copy.split)
+ self.assertTrue(x.lshape == x_copy.lshape)
+ self.assertTrue((x == x_copy).all().item())
+
+ # broadcasting shapes
+ x.resplit_(axis=0)
+ self.assert_array_equal(x[ht.array(k1, split=0), ht.array(1), 2], x_np[k1, 1, 2])
+ # test exception: broadcasting mismatching shapes
+ k2 = np.array([0, 2, 1])
+ with self.assertRaises(IndexError):
+ x[k1, k2, k3]
+
+ # more broadcasting
+ x_np = np.arange(12).reshape(4, 3)
+ rows = np.array([0, 3])
+ cols = np.array([0, 2])
+ x = ht.arange(12).reshape(4, 3)
+ x.resplit_(1)
+ x_np_indexed = x_np[rows[:, np.newaxis], cols]
+ x_indexed = x[ht.array(rows)[:, np.newaxis], cols]
+ self.assert_array_equal(x_indexed, x_np_indexed)
+ self.assertTrue(x_indexed.split == 1)
+
+ # 1d, split 0, advanced indexing with negative indices (fix #824)
+ x = ht.arange(10, 1, -1, split=0)
+ x_np = np.arange(10, 1, -1)
+ idx_np = np.array([3, 3, -3, 8])
+ idx = ht.array(idx_np)
+
+ x_adv_ind = x[idx]
+ x_np_adv_ind = x_np[idx_np]
+ self.assert_array_equal(x_adv_ind, x_np_adv_ind)
+
+ # 2d, split 0, multi-dimensional advanced indexing (fix #824)
+ x = ht.arange(10, 1, -1, split=0)
+ x_np = np.arange(10, 1, -1)
+ idx_np_2d = np.array([[1, 1], [2, 3]])
+ idx_2d = ht.array(idx_np_2d)
+
+ x_adv_ind_2d = x[idx_2d]
+ x_np_adv_ind_2d = x_np[idx_np_2d]
+ self.assert_array_equal(x_adv_ind_2d, x_np_adv_ind_2d)
+
+ # combining advanced and basic indexing
+ y_np = np.arange(35).reshape(5, 7)
+ y_np_indexed = y_np[np.array([0, 2, 4]), 1:3]
+ y = ht.array(y_np, split=1)
+ y_indexed = y[ht.array([0, 2, 4]), 1:3]
+ self.assert_array_equal(y_indexed, y_np_indexed)
+ self.assertTrue(y_indexed.split == 1)
+
+ x_np = np.arange(10 * 20 * 30).reshape(10, 20, 30)
+ x = ht.array(x_np, split=1)
+ ind_array = ht.random.randint(0, 20, (2, 3, 4), dtype=ht.int64)
+ ind_array_np = ind_array.numpy()
+ x_np_indexed = x_np[..., ind_array_np, :]
+ x_indexed = x[..., ind_array, :]
+ self.assert_array_equal(x_indexed, x_np_indexed)
+ self.assertTrue(x_indexed.split == 3)
+
+ # multi-array advanced indexing (consecutive dimensions)
+ arr_np = np.arange(4 * 5 * 6 * 7).reshape((4, 5, 6, 7))
+ arr = ht.array(arr_np, split=3)
+ a1_np = np.array([1, 2])
+ a2_np = np.array([3, 4])
+
+ a1 = ht.array(a1_np)
+ a2 = ht.array(a2_np)
+
+ res_consec_np = arr_np[:, a1_np, a2_np, :]
+ res_consec = arr[:, a1, a2, :]
+ self.assertEqual(res_consec.split, 2)
+ self.assertEqual(res_consec.gshape, (4, 2, 7))
+ self.assert_array_equal(res_consec, res_consec_np)
+
+ # multi-array advanced indexing (non-consecutive dimensions)
+ res_nonconsec_np = arr_np[a1_np, :, a2_np, :]
+ res_nonconsec = arr[a1, :, a2, :]
+ self.assert_array_equal(res_nonconsec, res_nonconsec_np)
+ self.assertEqual(res_nonconsec.split, 2)
+ self.assertEqual(res_nonconsec.gshape, (2, 5, 7))
+
+ def test_getitem_boolean_mask(self):
+ # boolean mask, local
+ arr = ht.arange(3 * 4 * 5).reshape(3, 4, 5)
+ np.random.seed(42)
+ mask = np.random.randint(0, 2, arr.shape, dtype=bool)
+ self.assertTrue((arr[mask].numpy() == arr.numpy()[mask]).all())
+
+ # boolean mask, distributed
+ arr_split0 = ht.array(arr, split=0)
+ mask_split0 = ht.array(mask, split=0)
+ self.assertTrue((arr_split0[mask_split0].numpy() == arr.numpy()[mask]).all())
+
+ arr_split1 = ht.array(arr, split=1)
+ mask_split1 = ht.array(mask, split=1)
+ self.assert_array_equal(arr_split1[mask_split1], arr.numpy()[mask])
+
+ arr_split2 = ht.array(arr, split=2)
+ mask_split2 = ht.array(mask, split=2)
+ self.assert_array_equal(arr_split2[mask_split2], arr.numpy()[mask])
+
+ # 0-D arrays indexed by Python booleans or 0-D boolean tensors
+ x_np = np.array(42)
+ x_ht = ht.array(42)
+
+ self.assert_array_equal(x_ht[False], x_np[False])
+ self.assert_array_equal(x_ht[True], x_np[True])
+ self.assert_array_equal(x_ht[ht.array(False)], x_np[np.array(False)])
+ self.assert_array_equal(x_ht[ht.array(True)], x_np[np.array(True)])
+
+ # boolean edge case
+ idx = ht.array([2, 0, 1], split=0)
+ mask = ht.array([True, False, True], split=0)
+ self.assertTrue((idx[mask] == ht.array([2, 1], dtype=idx.dtype, split=0)).all().item())
+
def test_int_cast(self):
# simple scalar tensor
a = ht.ones(1)
@@ -1248,380 +1583,407 @@ def test_resplit(self):
self.assertTrue(ht.all(t1_sub == res))
self.assertEqual(t1_sub.split, None)
- # 3D non-contiguous resplit testing (Column mayor ordering)
- torch_array = torch.arange(100, device=self.device.torch_device).reshape((10, 5, 2))
- heat_array = ht.array(torch_array, split=2, order="F")
- heat_array.resplit_(axis=1)
- res = np.arange(100).reshape(10, 5, 2)
- self.assertTrue(ht.array(res).device == heat_array.device)
- self.assertTrue(ht.all(heat_array == ht.array(res)))
- self.assertEqual(heat_array.split, 1)
-
- # 4D non-contiguous resplit testing (from transpose
- torch_array = torch.arange(5 * 4 * 3 * 6, device=self.device.torch_device).reshape(
- 5, 4, 3, 6
- )
- res = torch_array.cpu().numpy().transpose((3, 1, 2, 0))
- heat_array = ht.array(torch_array, split=2).transpose((3, 1, 2, 0))
- heat_array.resplit_(axis=1)
- self.assertTrue(ht.array(res).device == heat_array.device)
- self.assertTrue(ht.all(heat_array == ht.array(res)))
- self.assertEqual(heat_array.split, 1)
-
- def test_setitem_getitem(self):
+ # 3D non-contiguous resplit testing (Column major ordering)
+ torch_array = torch.arange(100, device=self.device.torch_device).reshape((10, 5, 2))
+ heat_array = ht.array(torch_array, split=2, order="F")
+ heat_array.resplit_(axis=1)
+ res = np.arange(100).reshape(10, 5, 2)
+ self.assertTrue(ht.array(res).device == heat_array.device)
+ self.assertTrue(ht.all(heat_array == ht.array(res)))
+ self.assertEqual(heat_array.split, 1)
+
+ # 4D non-contiguous resplit testing (from transpose
+ torch_array = torch.arange(5 * 4 * 3 * 6, device=self.device.torch_device).reshape(5, 4, 3, 6)
+ res = torch_array.cpu().numpy().transpose((3, 1, 2, 0))
+ heat_array = ht.array(torch_array, split=2).transpose((3, 1, 2, 0))
+ heat_array.resplit_(axis=1)
+ self.assertTrue(ht.array(res).device == heat_array.device)
+ self.assertTrue(ht.all(heat_array == ht.array(res)))
+ self.assertEqual(heat_array.split, 1)
+
+
+ def test_setitem_single_element(self):
+ # Single element indexing
+ # 1D, local
+ x = ht.zeros(10)
+ x[2] = 2
+ x[-2] = 8
+ self.assertTrue(x[2].item() == 2)
+ self.assertTrue(x[-2].item() == 8)
+ self.assertTrue(x[2].dtype == ht.float32)
+ # 1D, distributed
+ x = ht.zeros(10, split=0, dtype=ht.float64)
+ x[2] = 2
+ x[-2] = 8
+ self.assertTrue(x[2].item() == 2.0)
+ self.assertTrue(x[-2].item() == 8.0)
+ self.assertTrue(x[2].dtype == ht.float64)
+ self.assertTrue(x.split == 0)
+ # 2D, local
+ x = ht.zeros(10).reshape(2, 5)
+ x[0] = ht.arange(5)
+ self.assertTrue((x[0] == ht.arange(5)).all().item())
+ self.assertTrue(x[0].dtype == ht.float32)
+ # 2D, distributed
+ x_split0 = ht.zeros(10, split=0).reshape(2, 5)
+ x_split0[0] = ht.arange(5)
+ self.assertTrue((x_split0[0] == ht.arange(5, split=None)).all().item())
+ x_split1 = ht.zeros(10, split=0).reshape(2, 5, new_split=1)
+ x_split1[-2] = ht.arange(5)
+ self.assertTrue((x_split1[-2] == ht.arange(5, split=0)).all().item())
+ # 3D, distributed, split = 0
+ x_split0 = ht.zeros(27, split=0).reshape(3, 3, 3)
+ key = -2
+ x_split0[key] = ht.arange(3)
+ self.assertTrue((x_split0[key] == ht.arange(3, device=x_split0.device)).all().item())
+ self.assertTrue(x_split0[key].dtype == ht.float32)
+ self.assertTrue(x_split0.split == 0)
+ # 3D, distributed split, != 0
+ x_split2 = ht.zeros(27, dtype=ht.int64, split=0).reshape(3, 3, 3, new_split=2)
+ key = ht.array(2)
+ x_split2[key] = [6, 7, 8]
+ indexed_split2 = x_split2[key]
+ self.assertTrue((indexed_split2.numpy()[0] == np.array([6, 7, 8])).all())
+ self.assertTrue(indexed_split2.dtype == ht.int64)
+ self.assertTrue(x_split2.split == 2)
+
+ def test_setitem_slicing(self):
+ # Slicing and striding
+ x = ht.arange(20, split=0)
+ x[1:11:3] = ht.array([10, 40, 70, 100])
+ x_np = np.arange(20)
+ x_np[1:11:3] = np.array([10, 40, 70, 100])
+ self.assert_array_equal(x, x_np)
+ self.assertTrue(x.split == 0)
+
+ # 1-element slice along split axis
+ x = ht.arange(20).reshape(4, 5)
+ x.resplit_(axis=1)
+ x[:, 2:3] = ht.array([10, 40, 70, 100]).reshape(4, 1)
+ x_np = np.arange(20).reshape(4, 5)
+ x_np[:, 2:3] = np.array([10, 40, 70, 100]).reshape(4, 1)
+ self.assert_array_equal(x, x_np)
+ self.assertTrue(x.split == 1)
+ with self.assertRaises(ValueError):
+ x[:, 2:3] = ht.array([10, 40, 70, 100])
+
# tests for bug #825
a = ht.ones((102, 102), split=0)
setting = ht.zeros((100, 100), split=0)
a[1:-1, 1:-1] = setting
- self.assertTrue(ht.all(a[1:-1, 1:-1] == 0))
+ self.assertTrue(ht.all(a[1:-1, 1:-1] == 0).item())
a = ht.ones((102, 102), split=1)
setting = ht.zeros((30, 100), split=1)
a[-30:, 1:-1] = setting
- self.assertTrue(ht.all(a[-30:, 1:-1] == 0))
+ self.assertTrue(ht.all(a[-30:, 1:-1] == 0).item())
a = ht.ones((102, 102), split=1)
setting = ht.zeros((100, 100), split=1)
a[1:-1, 1:-1] = setting
- self.assertTrue(ht.all(a[1:-1, 1:-1] == 0))
+ self.assertTrue(ht.all(a[1:-1, 1:-1] == 0).item())
a = ht.ones((102, 102), split=1)
setting = ht.zeros((100, 20), split=1)
a[1:-1, :20] = setting
- self.assertTrue(ht.all(a[1:-1, :20] == 0))
+ self.assertTrue(ht.all(a[1:-1, :20] == 0).item())
- # tests for bug 730:
- a = ht.ones((10, 25, 30), split=1)
if a.comm.size > 1:
- self.assertEqual(a[0].split, 0)
- self.assertEqual(a[:, 0, :].split, None)
- self.assertEqual(a[:, :, 0].split, 1)
-
- # set and get single value
- a = ht.zeros((13, 5), split=0)
- # set value on one node
- a[10, np.array(0)] = 1
- self.assertEqual(a[10, 0], 1)
- self.assertEqual(a[10, 0].dtype, ht.float32)
-
- a = ht.zeros((13, 5), split=0)
- a[10] = 1
- b = a[torch.tensor(10)]
- self.assertTrue((b == 1).all())
- self.assertEqual(b.dtype, ht.float32)
- self.assertEqual(b.gshape, (5,))
-
- a = ht.zeros((13, 5), split=0)
- a[-1] = 1
- b = a[-1]
- self.assertTrue((b == 1).all())
- self.assertEqual(b.dtype, ht.float32)
- self.assertEqual(b.gshape, (5,))
-
- # slice in 1st dim only on 1 node
- a = ht.zeros((13, 5), split=0)
- a[1:4] = 1
- self.assertTrue((a[1:4] == 1).all())
- self.assertEqual(a[1:4].gshape, (3, 5))
- self.assertEqual(a[1:4].split, 0)
- self.assertEqual(a[1:4].dtype, ht.float32)
- if a.comm.size == 2:
- if a.comm.rank == 0:
- self.assertEqual(a[1:4].lshape, (3, 5))
- else:
- self.assertEqual(a[1:4].lshape, (0, 5))
-
- a = ht.zeros((13, 5), split=0)
- a[1:2] = 1
- self.assertTrue((a[1:2] == 1).all())
- self.assertEqual(a[1:2].gshape, (1, 5))
- self.assertEqual(a[1:2].split, 0)
- self.assertEqual(a[1:2].dtype, ht.float32)
- if a.comm.size == 2:
- if a.comm.rank == 0:
- self.assertEqual(a[1:2].lshape, (1, 5))
- else:
- self.assertEqual(a[1:2].lshape, (0, 5))
-
- # slice in 1st dim only on 1 node w/ singular second dim
- a = ht.zeros((13, 5), split=0)
- a[1:4, 1] = 1
- b = a[1:4, np.int64(1)]
- self.assertTrue((b == 1).all())
- self.assertEqual(b.gshape, (3,))
- self.assertEqual(b.split, 0)
- self.assertEqual(b.dtype, ht.float32)
- if a.comm.size == 2:
- if a.comm.rank == 0:
- self.assertEqual(b.lshape, (3,))
- else:
- self.assertEqual(b.lshape, (0,))
-
- # slice in 1st dim across both nodes (2 node case) w/ singular second dim
- a = ht.zeros((13, 5), split=0)
- a[1:11, 1] = 1
- self.assertTrue((a[1:11, 1] == 1).all())
- self.assertEqual(a[1:11, 1].gshape, (10,))
- self.assertEqual(a[1:11, torch.tensor(1)].split, 0)
- self.assertEqual(a[1:11, 1].dtype, ht.float32)
- if a.comm.size == 2:
- if a.comm.rank == 1:
- self.assertEqual(a[1:11, 1].lshape, (4,))
- if a.comm.rank == 0:
- self.assertEqual(a[1:11, 1].lshape, (6,))
-
- # slice in 1st dim across 1 node (2nd) w/ singular second dim
- c = ht.zeros((13, 5), split=0)
- c[8:12, ht.array(1)] = 1
- b = c[8:12, np.int64(1)]
- self.assertTrue((b == 1).all())
- self.assertEqual(b.gshape, (4,))
- self.assertEqual(b.split, 0)
- self.assertEqual(b.dtype, ht.float32)
- if a.comm.size == 2:
- if a.comm.rank == 1:
- self.assertEqual(b.lshape, (4,))
- if a.comm.rank == 0:
- self.assertEqual(b.lshape, (0,))
-
- # slice in both directions
- a = ht.zeros((13, 5), split=0)
- a[3:13, 2:5:2] = 1
- self.assertTrue((a[3:13, 2:5:2] == 1).all())
- self.assertEqual(a[3:13, 2:5:2].gshape, (10, 2))
- self.assertEqual(a[3:13, 2:5:2].split, 0)
- self.assertEqual(a[3:13, 2:5:2].dtype, ht.float32)
- if a.comm.size == 2:
- if a.comm.rank == 1:
- self.assertEqual(a[3:13, 2:5:2].lshape, (6, 2))
- if a.comm.rank == 0:
- self.assertEqual(a[3:13, 2:5:2].lshape, (4, 2))
-
- # setting with heat tensor
- a = ht.zeros((4, 5), split=0)
- if self.is_mps:
- a[1, 0:4] = ht.arange(4, dtype=a.dtype)
- else:
- a[1, 0:4] = ht.arange(4)
- # if a.comm.size == 2:
- for c, i in enumerate(range(4)):
- self.assertEqual(a[1, c], i)
+ with self.assertRaises(RuntimeError):
+ x = ht.ones((10, 10), split=0)
+ setting = ht.zeros((8, 8), split=1)
+ x[1:-1, 1:-1] = setting
- # setting with torch tensor
- a = ht.zeros((4, 5), split=0)
- if self.is_mps:
- a[1, 0:4] = torch.arange(4, dtype=a.larray.dtype, device=self.device.torch_device)
- else:
- a[1, 0:4] = torch.arange(4, device=self.device.torch_device)
- # if a.comm.size == 2:
- for c, i in enumerate(range(4)):
- self.assertEqual(a[1, c], i)
-
- ###################################################
- a = ht.zeros((13, 5), split=1)
- # # set value on one node
- a[10] = 1
- self.assertEqual(a[10].dtype, ht.float32)
- if a.comm.size == 2:
- if a.comm.rank == 0:
- self.assertEqual(a[10].lshape, (3,))
- if a.comm.rank == 1:
- self.assertEqual(a[10].lshape, (2,))
-
- a = ht.zeros((13, 5), split=1)
- # # set value on one node
- a[10, 0] = 1
- self.assertEqual(a[10, 0], 1)
- self.assertEqual(a[10, 0].dtype, ht.float32)
-
- # slice in 1st dim only on 1 node
- a = ht.zeros((13, 5), split=1)
- a[1:4] = 1
- self.assertTrue((a[1:4] == 1).all())
- self.assertEqual(a[1:4].gshape, (3, 5))
- self.assertEqual(a[1:4].split, 1)
- self.assertEqual(a[1:4].dtype, ht.float32)
- if a.comm.size == 2:
- if a.comm.rank == 0:
- self.assertEqual(a[1:4].lshape, (3, 3))
- if a.comm.rank == 1:
- self.assertEqual(a[1:4].lshape, (3, 2))
-
- # slice in 1st dim only on 1 node w/ singular second dim
- a = ht.zeros((13, 5), split=1)
- a[1:4, 1] = 1
- self.assertTrue((a[1:4, 1] == 1).all())
- self.assertEqual(a[1:4, 1].gshape, (3,))
- self.assertEqual(a[1:4, 1].split, None)
- self.assertEqual(a[1:4, 1].dtype, ht.float32)
- if a.comm.size == 2:
- if a.comm.rank == 0:
- self.assertEqual(a[1:4, 1].lshape, (3,))
- if a.comm.rank == 1:
- self.assertEqual(a[1:4, 1].lshape, (3,))
-
- # slice in 2st dim across both nodes (2 node case) w/ singular fist dim
- a = ht.zeros((13, 5), split=1)
- a[11, 1:5] = 1
- self.assertTrue((a[11, 1:5] == 1).all())
- self.assertEqual(a[11, 1:5].gshape, (4,))
- self.assertEqual(a[11, 1:5].split, 0)
- self.assertEqual(a[11, 1:5].dtype, ht.float32)
- if a.comm.size == 2:
- if a.comm.rank == 1:
- self.assertEqual(a[11, 1:5].lshape, (2,))
- if a.comm.rank == 0:
- self.assertEqual(a[11, 1:5].lshape, (2,))
-
- # slice in 1st dim across 1 node (2nd) w/ singular second dim
- a = ht.zeros((13, 5), split=1)
- a[8:12, 1] = 1
- self.assertTrue((a[8:12, 1] == 1).all())
- self.assertEqual(a[8:12, 1].gshape, (4,))
- self.assertEqual(a[8:12, 1].split, None)
- self.assertEqual(a[8:12, 1].dtype, ht.float32)
- if a.comm.size == 2:
- if a.comm.rank == 0:
- self.assertEqual(a[8:12, 1].lshape, (4,))
- if a.comm.rank == 1:
- self.assertEqual(a[8:12, 1].lshape, (4,))
-
- # slice in both directions
- a = ht.zeros((13, 5), split=1)
- a[3:13, 2::2] = 1
- self.assertTrue((a[3:13, 2:5:2] == 1).all())
- self.assertEqual(a[3:13, 2:5:2].gshape, (10, 2))
- self.assertEqual(a[3:13, 2:5:2].split, 1)
- self.assertEqual(a[3:13, 2:5:2].dtype, ht.float32)
- if a.comm.size == 2:
- if a.comm.rank == 1:
- self.assertEqual(a[3:13, 2:5:2].lshape, (10, 1))
- if a.comm.rank == 0:
- self.assertEqual(a[3:13, 2:5:2].lshape, (10, 1))
-
- a = ht.zeros((13, 5), split=1)
- a[..., 2::2] = 1
- self.assertTrue((a[:, 2:5:2] == 1).all())
- self.assertEqual(a[..., 2:5:2].gshape, (13, 2))
- self.assertEqual(a[..., 2:5:2].split, 1)
- self.assertEqual(a[..., 2:5:2].dtype, ht.float32)
- if a.comm.size == 2:
- if a.comm.rank == 1:
- self.assertEqual(a[..., 2:5:2].lshape, (13, 1))
- if a.comm.rank == 0:
- self.assertEqual(a[:, 2:5:2].lshape, (13, 1))
-
- # setting with heat tensor
- a = ht.zeros((4, 5), split=1)
- if self.is_mps:
- a[1, 0:4] = ht.arange(4, dtype=a.dtype)
- else:
- a[1, 0:4] = ht.arange(4)
- for c, i in enumerate(range(4)):
- b = a[1, c]
- if b.larray.numel() > 0:
- self.assertEqual(b.item(), i)
-
- # setting with torch tensor
- a = ht.zeros((4, 5), split=1)
- if a.device.torch_device.startswith("mps"):
- a[1, 0:4] = torch.arange(4, dtype=a.larray.dtype, device=self.device.torch_device)
- else:
- a[1, 0:4] = torch.arange(4, device=self.device.torch_device)
- for c, i in enumerate(range(4)):
- self.assertEqual(a[1, c], i)
-
- ####################################################
- a = ht.zeros((13, 5, 7), split=2)
- # # set value on one node
- a[10, :, :] = 1
- self.assertEqual(a[10, :, :].dtype, ht.float32)
- self.assertEqual(a[10, :, :].gshape, (5, 7))
- if a.comm.size == 2:
- if a.comm.rank == 0:
- self.assertEqual(a[10, :, :].lshape, (5, 4))
- if a.comm.rank == 1:
- self.assertEqual(a[10, :, :].lshape, (5, 3))
-
- a = ht.zeros((13, 5, 7), split=2)
- # # set value on one node
- a[10, ...] = 1
- self.assertEqual(a[10, ...].dtype, ht.float32)
- self.assertEqual(a[10, ...].gshape, (5, 7))
- if a.comm.size == 2:
- if a.comm.rank == 0:
- self.assertEqual(a[10, ...].lshape, (5, 4))
- if a.comm.rank == 1:
- self.assertEqual(a[10, ...].lshape, (5, 3))
-
- a = ht.zeros((13, 5, 8), split=2)
- # # set value on one node
- a[10, 0, 0] = 1
- self.assertEqual(a[10, 0, 0], 1)
- self.assertEqual(a[10, 0, 0].dtype, ht.float32)
-
- # # slice in 1st dim only on 1 node
- a = ht.zeros((13, 5, 7), split=2)
- a[1:4] = 1
- self.assertTrue((a[1:4] == 1).all())
- self.assertEqual(a[1:4].gshape, (3, 5, 7))
- self.assertEqual(a[1:4].split, 2)
- self.assertEqual(a[1:4].dtype, ht.float32)
- if a.comm.size == 2:
- if a.comm.rank == 0:
- self.assertEqual(a[1:4].lshape, (3, 5, 4))
- if a.comm.rank == 1:
- self.assertEqual(a[1:4].lshape, (3, 5, 3))
-
- # slice in 1st dim only on 1 node w/ singular second dim
- a = ht.zeros((13, 5, 7), split=2)
- a[1:4, 1, :] = 1
- self.assertTrue((a[1:4, 1, :] == 1).all())
- self.assertEqual(a[1:4, 1, :].gshape, (3, 7))
- if a.comm.size == 2:
- self.assertEqual(a[1:4, 1, :].split, 1)
- self.assertEqual(a[1:4, 1, :].dtype, ht.float32)
- if a.comm.rank == 0:
- self.assertEqual(a[1:4, 1, :].lshape, (3, 4))
- if a.comm.rank == 1:
- self.assertEqual(a[1:4, 1, :].lshape, (3, 3))
-
- # slice in both directions
- a = ht.zeros((13, 5, 7), split=2)
- a[3:13, 2:5:2, 1:7:3] = 1
- self.assertTrue((a[3:13, 2:5:2, 1:7:3] == 1).all())
- self.assertEqual(a[3:13, 2:5:2, 1:7:3].split, 2)
- self.assertEqual(a[3:13, 2:5:2, 1:7:3].dtype, ht.float32)
- self.assertEqual(a[3:13, 2:5:2, 1:7:3].gshape, (10, 2, 2))
- if a.comm.size == 2:
- out = ht.ones((4, 5, 5), split=1)
- self.assertEqual(out[0].gshape, (5, 5))
- if a.comm.rank == 1:
- self.assertEqual(a[3:13, 2:5:2, 1:7:3].lshape, (10, 2, 1))
- self.assertEqual(out[0].lshape, (2, 5))
- if a.comm.rank == 0:
- self.assertEqual(a[3:13, 2:5:2, 1:7:3].lshape, (10, 2, 1))
- self.assertEqual(out[0].lshape, (3, 5))
-
- a = ht.ones((4, 5), split=0).tril()
- a[0] = [6, 6, 6, 6, 6]
- self.assertTrue((a[0] == 6).all())
-
- a = ht.ones((4, 5), split=0).tril()
- a[0] = (6, 6, 6, 6, 6)
- self.assertTrue((a[0] == 6).all())
-
- a = ht.ones((4, 5), split=0).tril()
- a[0] = np.array([6, 6, 6, 6, 6])
- self.assertTrue((a[0] == 6).all())
-
- a = ht.ones((4, 5), split=0).tril()
- a[0] = ht.array([6, 6, 6, 6, 6])
- self.assertTrue((a[ht.array((0,))] == 6).all())
-
- a = ht.ones((4, 5), split=0).tril()
- a[0] = ht.array([6, 6, 6, 6, 6])
- self.assertTrue((a[ht.array((0,))] == 6).all())
+ def test_setitem_slicing_negative_step(self):
+ # slicing with negative step along split axis 0
+ # assign different dtype
+ shape = (20, 4, 3)
+ x_3d = ht.arange(20 * 4 * 3, split=0).reshape(shape)
+ value = ht.random.randn(8, 2)
+ x_3d[17:2:-2, :2, ht.array(1)] = value
+ x_3d_sliced = x_3d[17:2:-2, :2, ht.array(1)]
+ self.assertTrue(ht.allclose(x_3d_sliced, value.astype(x_3d.dtype)))
+ self.assertTrue(x_3d_sliced.dtype == x_3d.dtype)
+
+ # slicing with negative step along split 1
+ shape = (4, 20, 3)
+ x_3d = ht.arange(20 * 4 * 3, dtype=ht.float32).reshape(shape)
+ x_3d.resplit_(axis=1)
+ key = (slice(None, 2), slice(17, 2, -2), 1)
+ value = ht.random.randn(2, 8)
+ x_3d[key] = value
+ x_3d_sliced = x_3d[key]
+ self.assertTrue(ht.allclose(x_3d_sliced, value.astype(x_3d.dtype)))
+ self.assertTrue(x_3d_sliced.dtype == x_3d.dtype)
+
+ # slicing with negative step along split 2 and loss of axis < split
+ shape = (4, 3, 20)
+ x_3d = ht.arange(20 * 4 * 3, dtype=ht.float64).reshape(shape)
+ x_3d.resplit_(axis=2)
+ key = (slice(None, 2), 1, slice(17, 10, -2))
+ value = ht.random.randn(2, 4)
+ x_3d[key] = value
+ x_3d_sliced = x_3d[key]
+ self.assertTrue(ht.allclose(x_3d_sliced, value.astype(x_3d.dtype)))
+ self.assertTrue(x_3d_sliced.dtype == x_3d.dtype)
+
+ # slicing with negative step along split 2 and loss of all axes but split
+ shape = (4, 3, 20)
+ x_3d = ht.arange(20 * 4 * 3).reshape(shape)
+ x_3d.resplit_(axis=2)
+ key = (0, 1, slice(17, 13, -1))
+ value = ht.random.randint(
+ 0,
+ 5,
+ (
+ 1,
+ 4,
+ ),
+ split=1,
+ )
+ x_3d[key] = value
+ x_3d_sliced = x_3d[key]
+ self.assertTrue(ht.allclose(x_3d_sliced, value.squeeze(0).astype(x_3d.dtype)))
+ self.assertTrue(x_3d_sliced.dtype == x_3d.dtype)
+
+ def test_setitem_dimensional_indexing(self):
+ # ellipsis
+ x = ht.array([[[1], [2], [3]], [[4], [5], [6]]])
+ # local
+ value = x.squeeze() + 7
+ x[..., 0] = value
+ self.assertTrue(ht.all(x[..., 0] == value).item())
+ value -= 7
+ x[:, :, 0] = value
+ self.assertTrue(ht.all(x[:, :, 0] == value).item())
+
+ # distributed
+ x.resplit_(axis=1)
+ value *= 2
+ x[..., 0] = value
+ x_ellipsis = x[..., 0]
+ self.assertTrue(ht.all(x_ellipsis == value).item())
+ value += 2
+ x[:, :, 0] = value
+ self.assertTrue(ht.all(x[:, :, 0] == value).item())
+ self.assertTrue(x_ellipsis.split == 1)
+
+ # newaxis: local, w. broadcasting and different dtype
+ x = ht.array([[[1], [2], [3]], [[4], [5], [6]]])
+ value = ht.array([10.0, 20.0]).reshape(2, 1)
+ x[:, None, :2, :] = value
+ x_newaxis = x[:, None, :2, :]
+ self.assertTrue(ht.all(x_newaxis == value).item())
+ value += 2
+ x[:, None, :2, :] = value
+ self.assertTrue(ht.all(x[:, None, :2, :] == value).item())
+ self.assertTrue(x[:, None, :2, :].dtype == x.dtype)
+
+ # newaxis: distributed w. broadcasting and different dtype
+ x.resplit_(axis=1)
+ value = ht.array([30.0, 40.0]).reshape(1, 2, 1)
+ x[:, np.newaxis, :2, :] = value
+ x_newaxis = x[:, np.newaxis, :2, :]
+ self.assertTrue(ht.all(x_newaxis == value).item())
+ value += 2
+ x[:, None, :2, :] = value
+ x_none = x[:, None, :2, :]
+ self.assertTrue(ht.all(x_none == value).item())
+ self.assertTrue(x_none.dtype == x.dtype)
+
+ # distributed value
+ x = ht.arange(6).reshape(1, 1, 2, 3)
+ x.resplit_(axis=-1)
+ value = ht.arange(3).reshape(1, 3)
+ value.resplit_(axis=1)
+ x[..., 0, :] = value
+ self.assertTrue(ht.all(x[..., 0, :] == value).item())
+
+ # test multiple ellipses rejection
+ a = ht.ones((5, 5))
+ with self.assertRaises(ValueError):
+ a[..., ...] = 0
+
+ def test_setitem_advanced_indexing(self):
+ # "x[(1, 2, 3),] is fundamentally different from x[(1, 2, 3)]", cf. numpy docs
+
+ x = ht.arange(60, split=0).reshape(5, 3, 4)
+ value = 99.0
+ x[(1, 2, 3)] = value
+ indexed_x = x[(1, 2, 3)]
+ self.assertTrue((indexed_x == value).item())
+ self.assertTrue(indexed_x.dtype == x.dtype)
+ x[(1, 2, 3),] = value
+ adv_indexed_x = x[(1, 2, 3),]
+ self.assertTrue(ht.all(adv_indexed_x == value).item())
+ self.assertTrue(adv_indexed_x.dtype == x.dtype)
+
+ # 1d
+ x = ht.arange(10, 1, -1, split=0)
+ value = ht.arange(4)
+ x[ht.array([3, 2, 1, 8])] = value
+ x_adv_ind = x[np.array([3, 2, 1, 8])]
+ self.assertTrue(ht.all(x_adv_ind == value).item())
+ self.assertTrue(x_adv_ind.dtype == x.dtype)
+
+ # 1d, split 0, advanced indexing with a local DNDarray
+ x = ht.arange(10, 1, -1, split=0)
+ x_np = np.arange(10, 1, -1)
+ idx_np = np.array([3, 3, 1, 8])
+ idx = ht.array(idx_np, split=None) # Explicitly local DNDarray
+ vals_np = np.arange(4)
+ vals = ht.array(vals_np, split=None)
+ x[idx] = vals
+ x_np[idx_np] = vals_np
+ self.assertTrue(ht.all(x == ht.array(x_np, split=0)).item())
+
+ # 2d, split 0, single 1d tensor unordered advanced indexing
+ arr = ht.zeros((10, 5), dtype=ht.float32, split=0)
+ idx_np = np.array([7, 2, 8, 1])
+ idx = ht.array(idx_np, split=0)
+
+ vals_np = np.arange(20, dtype=np.float32).reshape(4, 5)
+ vals = ht.array(vals_np, split=0)
+
+ arr[idx] = vals
+
+ arr_np = np.zeros((10, 5), dtype=np.float32)
+ arr_np[idx_np] = vals_np
+ self.assertTrue((arr == ht.array(arr_np, split=0)).all().item())
+
+ # 3d, split 0, non-unique, non-ordered key along split axis, key mask-like
+ x = ht.arange(60, split=0).reshape(5, 3, 4)
+ k1 = np.array([0, 4, 1, 0])
+ k2 = np.array([0, 2, 1, 0])
+ k3 = np.array([1, 2, 3, 1])
+ value = ht.array([99, 98, 97, 96], split=0)
+ x[k1, k2, k3] = value
+ self.assertTrue((x[k1, k2, k3] == ht.array([96, 98, 97, 96], split=0)).all().item())
+
+ # advanced indexing on non-consecutive dimensions, split dimension will be lost
+ x = ht.arange(60, split=0).reshape(5, 3, 4, new_split=1)
+ x_copy = x.copy()
+ k1 = np.array([0, 4, 1, 2])
+ k2 = 0
+ k3 = np.array([1, 2, 3, 1])
+ key = (k1, k2, k3)
+ value = ht.array([99, 98, 97, 96])
+ x[key] = value
+ self.assertTrue((x[key] == ht.array([99, 98, 97, 96])).all().item())
+ # check that x is unchanged after internal manipulation
+ self.assertTrue(x.shape == x_copy.shape)
+ self.assertTrue(x.split == x_copy.split)
+ self.assertTrue(x.lshape == x_copy.lshape)
+
+ # broadcasting shapes
+ x.resplit_(axis=0)
+ key = (ht.array(k1, split=0), ht.array(1), 2)
+ value = ht.array([99, 98, 97, 96], split=0)
+ x[key] = value
+ self.assertTrue((x[key] == value).all().item())
+ # test exception: broadcasting mismatching shapes
+ k2 = np.array([0, 2, 1])
+ with self.assertRaises(IndexError):
+ x[k1, k2, k3] = value
+
+ # more broadcasting
+ x = ht.arange(12).reshape(4, 3)
+ x.resplit_(1)
+ rows = np.array([0, 3])
+ cols = np.array([0, 2])
+ key = (ht.array(rows)[:, np.newaxis], cols)
+ value = ht.array([[99, 98], [97, 96]], split=1)
+ x[key] = value
+ self.assertTrue((x[key] == value).all().item())
+ if x.comm.size > 1:
+ with self.assertRaises(RuntimeError):
+ value = ht.array([[99, 98], [97, 96]], split=0)
+ x[key] = value
+
+ # 1d, split 0, advanced indexing assignment with negative indices
+ x = ht.arange(10, 1, -1, split=0)
+ x_np = np.arange(10, 1, -1)
+ idx_np = np.array([3, 3, -3, 8])
+ idx = ht.array(idx_np)
+
+ vals_np = np.array([100, 101, 102, 103])
+ vals = ht.array(vals_np)
+
+ x[idx] = vals
+ x_np[idx_np] = vals_np
+ self.assert_array_equal(x, x_np)
+
+ # 2d, split 0, multi-dimensional advanced indexing assignment
+ x = ht.arange(10, 1, -1, split=0)
+ x_np = np.arange(10, 1, -1)
+ idx_np_2d = np.array([[1, 1], [2, 3]])
+ idx_2d = ht.array(idx_np_2d)
+
+ vals_np_2d = np.array([[200, 201], [202, 203]])
+ vals_2d = ht.array(vals_np_2d)
+
+ x[idx_2d] = vals_2d
+ x_np[idx_np_2d] = vals_np_2d
+ self.assert_array_equal(x, x_np)
+
+ # combining advanced and basic indexing
+
+ y = ht.arange(35).reshape(5, 7)
+ y.resplit_(1)
+ y_copy = y.copy()
+ # assign non-distributed value
+ value = ht.arange(6).reshape(3, 2)
+ y[ht.array([0, 2, 4]), 1:3] = value
+ self.assertTrue((y[ht.array([0, 2, 4]), 1:3] == value).all().item())
+ # assign distributed value
+ value.resplit_(1)
+ y_copy[ht.array([0, 2, 4]), 1:3] = value
+ self.assertTrue((y_copy[ht.array([0, 2, 4]), 1:3] == value).all().item())
+
+
+ x = ht.arange(10 * 20 * 30).reshape(10, 20, 30)
+ x.resplit_(1)
+ ind_array = ht.array(
+ torch.tensor(
+ [
+ [[11, 10, 3, 2], [13, 10, 0, 4], [9, 3, 2, 0]],
+ [[6, 10, 3, 8], [16, 10, 12, 9], [10, 18, 6, 15]],
+ ]
+ ),
+ dtype=ht.int64,
+ )
+ value = ht.ones((1, 2, 3, 4, 1))
+ x[..., ind_array, :] = value
+ self.assertTrue((x[..., ind_array, :] == value).all().item())
+
+ def test_setitem_boolean_mask(self):
+ # boolean mask, local
+ arr = ht.arange(3 * 4 * 5).reshape(3, 4, 5)
+ np.random.seed(42)
+ mask = np.random.randint(0, 2, arr.shape, dtype=bool)
+ value = 99.0
+ arr[mask] = value
+ self.assertTrue((arr[mask] == value).all().item())
+ self.assertTrue(arr[mask].dtype == arr.dtype)
+ value = ht.ones_like(arr)
+ arr[mask] = value[mask]
+ self.assertTrue((arr[mask] == value[mask]).all().item())
+
+ # boolean mask, distributed, non-distributed `value`
+ arr_split0 = ht.array(arr, split=0)
+ mask_split0 = ht.array(mask, split=0)
+ arr_split0[mask_split0] = value[mask]
+ indexed_arr = arr_split0[mask_split0]
+ indexed_arr.balance_()
+ self.assertTrue((indexed_arr == value[mask]).all().item())
+ arr_split1 = ht.array(arr, split=1)
+ mask_split1 = ht.array(mask, split=1)
+ arr_split1[mask_split1] = value[mask]
+ self.assertTrue((arr_split1[mask_split1] == value[mask]).all().item())
+ arr_split2 = ht.array(arr, split=2)
+ mask_split2 = ht.array(mask, split=2)
+ arr_split2[mask_split2] = value[mask]
+ self.assertTrue((arr_split2[mask_split2] == value[mask]).all().item())
# ======================= indexing with bools =================================
split = None
@@ -1697,29 +2059,6 @@ def test_setitem_getitem(self):
self.assertTrue(np.all(arr.numpy() == np_arr))
self.assertTrue(ht.all(arr[ht_key] == 10.0))
- with self.assertRaises(ValueError):
- a[..., ...]
- with self.assertRaises(ValueError):
- a[..., ...] = 1
- if a.comm.size > 1:
- with self.assertRaises(ValueError):
- x = ht.ones((10, 10), split=0)
- setting = ht.zeros((8, 8), split=1)
- x[1:-1, 1:-1] = setting
-
- for split in [None, 0, 1, 2]:
- for new_dim in [0, 1, 2]:
- for add in [np.newaxis, None]:
- arr = ht.ones((4, 3, 2), split=split, dtype=ht.int32)
- check = torch.ones((4, 3, 2), dtype=torch.int32)
- idx = [slice(None), slice(None), slice(None)]
- idx[new_dim] = add
- idx = tuple(idx)
- arr = arr[idx]
- check = check[idx]
- self.assertTrue(arr.shape == check.shape)
- self.assertTrue(arr.lshape[new_dim] == 1)
-
def test_size_gnumel(self):
a = ht.zeros((10, 10, 10), split=None)
self.assertEqual(a.size, 10 * 10 * 10)
@@ -1914,6 +2253,7 @@ def test_torch_proxy(self):
dndarray_proxy.storage().size() * dndarray_proxy.storage().element_size()
)
self.assertTrue(dndarray_proxy_nbytes == 1)
+ self.assertTrue(dndarray_proxy.names.index("split") == 1)
def test_torch_function(self):
arr = ht.array([1, 2, 3, 4])
@@ -1930,3 +2270,153 @@ def test_xor(self):
self.assertTrue(
ht.equal(int16_tensor ^ int16_vector, ht.bitwise_xor(int16_tensor, int16_vector))
)
+
+ def test_getitem_boolean_fewer_dims(self):
+ # Test case: 2D array, 1D boolean mask (selects rows)
+ # NumPy behavior: x_2D[bool_1D] selects entire rows
+ arr_np = np.arange(20).reshape((10, 2))
+ mask_np = np.array([True, False, True, False, True, False, True, False, True, False])
+ result_np = arr_np[mask_np] # Shape (5, 2)
+
+ # Case 1: split=None (local)
+ arr_ht = ht.array(arr_np, split=None)
+ mask_ht = ht.array(mask_np, split=None)
+ result_ht = arr_ht[mask_ht]
+ self.assert_array_equal(result_ht, result_np)
+ self.assertEqual(result_ht.split, None)
+ self.assertEqual(result_ht.gshape, (5, 2))
+
+ # Case 2: split=0 (split on the indexed dimension)
+ arr_ht_s0 = ht.array(arr_np, split=0)
+ mask_ht_s0 = ht.array(mask_np, split=0)
+
+ result_ht_s0 = arr_ht_s0[mask_ht_s0]
+
+ self.assert_array_equal(result_ht_s0, result_np)
+ self.assertEqual(result_ht_s0.split, 0)
+ self.assertEqual(result_ht_s0.gshape, (5, 2))
+
+ # Case 3: split=1 (split on a non-indexed dimension)
+ arr_ht_s1 = ht.array(arr_np, split=1)
+ # Mask can be local or split=0, test local (None) for broadcasting
+ mask_ht_sNone = ht.array(mask_np, split=None)
+ result_ht_s1 = arr_ht_s1[mask_ht_sNone]
+ self.assert_array_equal(result_ht_s1, result_np)
+ self.assertEqual(result_ht_s1.split, 1)
+ self.assertEqual(result_ht_s1.gshape, (5, 2))
+
+ # Case 4: 3D array, 2D boolean mask
+ arr_np_3d = np.arange(30).reshape((2, 3, 5))
+ mask_np_2d = np.array([[True, True, False], [False, True, True]])
+ result_np_3d = arr_np_3d[mask_np_2d] # Shape (4, 5)
+
+ # Test split=None
+ arr_ht_3d = ht.array(arr_np_3d, split=None)
+ mask_ht_2d = ht.array(mask_np_2d, split=None)
+ result_ht_3d = arr_ht_3d[mask_ht_2d]
+ self.assert_array_equal(result_ht_3d, result_np_3d)
+ self.assertEqual(result_ht_3d.gshape, (4, 5))
+
+ # Test split=2 (split on the non-indexed dimension)
+ arr_ht_3d_s2 = ht.array(arr_np_3d, split=2)
+ mask_ht_2d_sNone = ht.array(mask_np_2d, split=None) # Broadcast mask
+ result_ht_3d_s2 = arr_ht_3d_s2[mask_ht_2d_sNone]
+ self.assert_array_equal(result_ht_3d_s2, result_np_3d)
+ self.assertEqual(result_ht_3d_s2.gshape, (4, 5))
+ self.assertEqual(result_ht_3d_s2.split, 1) # New split axis (originally 2, 2 dims removed)
+
+ def test_setitem_boolean_fewer_dims(self):
+ # Test case: 2D array, 1D boolean mask (selects rows)
+ arr_np = np.arange(20).reshape((10, 2))
+ mask_np = np.array([True, False, True, False, True, False, True, False, True, False])
+ value = 99
+ arr_np_set = arr_np.copy()
+ arr_np_set[mask_np] = value
+
+ # Case 1: split=None (local)
+ arr_ht = ht.array(arr_np, split=None)
+ mask_ht = ht.array(mask_np, split=None)
+ arr_ht[mask_ht] = value
+ self.assert_array_equal(arr_ht, arr_np_set)
+
+ # Case 2: split=0 (split on the indexed dimension)
+ arr_ht_s0 = ht.array(arr_np, split=0)
+ mask_ht_s0 = ht.array(mask_np, split=0)
+ arr_ht_s0[mask_ht_s0] = value
+ self.assert_array_equal(arr_ht_s0, arr_np_set)
+
+ # Case 3: split=1 (split on a non-indexed dimension)
+ arr_ht_s1 = ht.array(arr_np, split=1)
+ mask_ht_sNone = ht.array(mask_np, split=None)
+ arr_ht_s1[mask_ht_sNone] = value
+ self.assert_array_equal(arr_ht_s1, arr_np_set)
+
+ def test_getitem_edge_cases(self):
+ # Test edge cases from NumPy docs
+
+ # Case 1: 0-D (Scalar) DNDarray
+ x_ht_0d = ht.array(10)
+ self.assertEqual(x_ht_0d.ndim, 0)
+ result_0d = x_ht_0d[()]
+ # NumPy returns a scalar, heat returns a 0-D tensor
+ self.assertEqual(result_0d.ndim, 0)
+ self.assertEqual(result_0d.item(), 10)
+
+ # Case 2: N-D local DNDarray
+ arr_np = np.arange(10).reshape((5, 2))
+ arr_ht_local = ht.array(arr_np, split=None)
+
+ # Test [...]
+ result_ellipsis = arr_ht_local[...]
+ self.assert_array_equal(result_ellipsis, arr_np)
+ self.assertIs(result_ellipsis.larray, arr_ht_local.larray) # Check for view
+
+ # Test [()]
+ result_empty_tuple = arr_ht_local[()]
+ self.assert_array_equal(result_empty_tuple, arr_np)
+ self.assertIs(result_empty_tuple.larray, arr_ht_local.larray) # Check for view
+
+ # Case 3: N-D split DNDarray
+ arr_ht_split = ht.array(arr_np, split=0)
+
+ # Test [...]
+ result_split_ellipsis = arr_ht_split[...]
+ self.assert_array_equal(result_split_ellipsis, arr_np)
+ self.assertEqual(result_split_ellipsis.split, 0)
+ self.assertIs(result_split_ellipsis.larray, arr_ht_split.larray) # Check for view
+
+ # Test [()]
+ result_split_empty_tuple = arr_ht_split[()]
+ self.assert_array_equal(result_split_empty_tuple, arr_np)
+ self.assertEqual(result_split_empty_tuple.split, 0)
+ self.assertIs(result_split_empty_tuple.larray, arr_ht_split.larray) # Check for view
+
+ def test_setitem_edge_cases(self):
+ # Test edge cases from NumPy docs
+
+ # Case 1: 0-D (Scalar) DNDarray
+ x_ht_0d = ht.array(10)
+ x_ht_0d[()] = 99
+ self.assertEqual(x_ht_0d.item(), 99)
+
+ # Case 2: N-D local DNDarray
+ arr_ht_local = ht.ones((5, 2), split=None)
+
+ # Test [...]
+ arr_ht_local[...] = 99
+ self.assertTrue(ht.all(arr_ht_local == 99).item())
+
+ # Test [()]
+ arr_ht_local[()] = 100
+ self.assertTrue(ht.all(arr_ht_local == 100).item())
+
+ # Case 3: N-D split DNDarray
+ arr_ht_split = ht.ones((5, 2), split=0)
+
+ # Test [...]
+ arr_ht_split[...] = 99
+ self.assertTrue(ht.all(arr_ht_split == 99).item())
+
+ # Test [()]
+ arr_ht_split[()] = 100
+ self.assertTrue(ht.all(arr_ht_split == 100).item())
diff --git a/tests/core/test_indexing.py b/tests/core/test_indexing.py
index 61dda3fa4f..ada5e55220 100644
--- a/tests/core/test_indexing.py
+++ b/tests/core/test_indexing.py
@@ -1,6 +1,7 @@
import heat as ht
from heat.testing.basic_test import TestCase
+import torch
class TestIndexing(TestCase):
def test_nonzero(self):
@@ -9,18 +10,18 @@ def test_nonzero(self):
a = ht.array([[1, 2, 3], [4, 5, 2], [7, 8, 9]], split=None)
cond = a > 3
nz = ht.nonzero(cond)
- self.assertEqual(nz.gshape, (5, 2))
- self.assertEqual(nz.dtype, ht.int64)
- self.assertEqual(nz.split, None)
+ self.assertEqual(len(nz), 2)
+ self.assertEqual(len(nz[0]), 5)
+ self.assertEqual(nz[0].dtype, ht.int64)
# split
a = ht.array([[1, 2, 3], [4, 5, 6], [7, 8, 9]], split=1)
cond = a > 3
nz = cond.nonzero()
- self.assertEqual(nz.gshape, (6, 2))
- self.assertEqual(nz.dtype, ht.int64)
- self.assertEqual(nz.split, 0)
- a[nz] = 10.0
+ self.assertEqual(len(nz), 2)
+ self.assertEqual(len(nz[0]), 6)
+ self.assertEqual(nz[0].dtype, ht.int64)
+ a[nz] = 10
self.assertEqual(ht.all(a[nz] == 10), 1)
# edge case: single non-zero element
@@ -28,11 +29,23 @@ def test_nonzero(self):
a = ht.zeros((4, 3), dtype=ht.bool, split=split)
a[1, 2] = True
nz = ht.indexing.nonzero(a)
- a.resplit_(None)
- nz.resplit_(None)
- self.assertEqual(nz.gshape, (1, 2))
self.assertTrue(ht.allclose(a[nz], a[a]))
+ a.comm.Barrier()
+
+ # as_tuple = False (torch-style output)
+ a = ht.array([[1, 0, 0], [0, 4, 1], [0, 6, 0]], split=1)
+ nz = ht.nonzero(a, as_tuple=False)
+ self.assertEqual(nz.gshape, (4, 2))
+ self.assertEqual(nz.dtype, ht.int64)
+ self.assertEqual(nz.split, 0)
+ t_a = a.resplit_(None).larray
+ t_nz = torch.nonzero(t_a, as_tuple=False)
+ self.assertTrue(ht.equal(nz, ht.array(t_nz)))
+ # attribute error
+ a = a.numpy()
+ with self.assertRaises(TypeError):
+ ht.nonzero(a)
def test_where(self):
# cases to test
@@ -40,16 +53,18 @@ def test_where(self):
a = ht.array([[1, 2, 3], [4, 5, 6], [7, 8, 9]], split=None)
cond = a > 3
wh = ht.where(cond)
- self.assertEqual(wh.gshape, (6, 2))
- self.assertEqual(wh.dtype, ht.int64)
- self.assertEqual(wh.split, None)
+ self.assertEqual(len(wh), 2)
+ self.assertEqual(wh[0].gshape[0], 6)
+ self.assertEqual(wh[0].dtype, ht.int64)
+ self.assertEqual(wh[0].split, None)
# split
a = ht.array([[1, 2, 3], [4, 5, 6], [7, 8, 9]], split=1)
cond = a > 3
wh = ht.where(cond)
- self.assertEqual(wh.gshape, (6, 2))
- self.assertEqual(wh.dtype, ht.int64)
- self.assertEqual(wh.split, 0)
+ self.assertEqual(len(wh), 2)
+ self.assertEqual(wh[0].gshape, (6,))
+ self.assertEqual(wh[0].dtype, ht.int64)
+ self.assertEqual(wh[0].split, 0)
# not split cond
a = ht.array([[0.0, 1.0, 2.0], [0.0, 2.0, 4.0], [0.0, 3.0, 6.0]], split=None)