"""Provides Heat's core data structure, the DNDarray, a distributed n-dimensional array"""
from __future__ import annotations
import bisect
import math
import numpy as np
import torch
import warnings
from mpi4py import MPI
from pathlib import Path
from enum import Enum
from typing import Any, Union, TypeVar
from collections.abc import Iterable
warnings.simplefilter("always", ResourceWarning)
# NOTE: heat module imports need to be placed at the very end of the file to avoid cyclic dependencies
__all__ = ["DNDarray"]
Communication = TypeVar("Communication")
# Type aliases
Index = Union[int, slice, type(...), None, torch.Tensor, np.ndarray, "DNDarray"]
Indexer = Union[Index, tuple[Index, ...], list[Index]]
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", "descending_slice", "distr_mask", "local_mask", "local", "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
def _unwrap_local_key(key: Any, device: torch.device | None = None) -> Any:
"""
Recursively unwrap local DNDarray or numpy array keys into torch-compatible indexers on correct device.
"""
if isinstance(key, DNDarray):
if key.is_distributed():
raise TypeError("Cannot use distributed DNDarray for local fast-path indexing")
return key.larray.item() if key.ndim == 0 else key.larray
if isinstance(key, np.ndarray):
return torch.from_numpy(key).to(device=device)
if isinstance(key, tuple):
if len(key) == 1 and _is_boolean_array(key[0]):
return _unwrap_local_key(key[0], device=device)
else:
return tuple(_unwrap_local_key(k, device=device) for k in key)
if isinstance(key, list):
return [_unwrap_local_key(k, device=device) for k in key]
return key
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``.
"""
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()
root = bisect.bisect_right(displs, key) - 1
# correct key for rank-specific displacement
if return_local_indices and arr.comm.rank == root:
key -= displs[root]
else:
root = None
return key, root
def _is_boolean_array(k: Any) -> bool:
"""Return True if k is a boolean or uint8 array/tensor of any dimension."""
return hasattr(k, "dtype") and k.dtype in (
ht_bool,
ht_uint8,
torch.bool,
torch.uint8,
np.bool_,
np.uint8,
)
def _is_boolean_scalar(k: Any) -> bool:
"""Return True if k is a python bool or a 0-D boolean array/tensor."""
return isinstance(k, bool) or (_is_boolean_array(k) and getattr(k, "ndim", 0) == 0)
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 not torch.is_tensor(rhs_in) or 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 = [
torch.where(t < 0, t + int(target_shape[d]), t)
.to(device=device, dtype=torch.int64)
.reshape(-1)
for d, t in enumerate(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]
# Determine sorting order (stable sort preserves original order)
order = torch.argsort(lin, stable=True)
pos = None
lin_s = lin[order]
# Fast path: check adjacent elements in sorted order
# If all adjacent elements are distinct, there are no duplicates
if (lin_s[1:] != lin_s[:-1]).all():
return key_in, rhs_in
if pos is None:
pos = torch.arange(n, device=device, dtype=torch.int64)
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 _is_scalar_index(k: Any) -> bool:
"""Return True if k is a non-boolean scalar or 0-D array indexer."""
return not _is_boolean_scalar(k) and (np.isscalar(k) or getattr(k, "ndim", 1) == 0)
def _normalize_key(key: Indexer, device: torch.device) -> tuple[Any, ...]:
"""
Standardize the non-DNDarray coordinate indices to PyTorch-friendly key items.
Returns a tuple of normalized key items.
"""
# Normalize top-level container to a list
if isinstance(key, tuple):
key_list = list(key)
elif isinstance(key, list):
# Could be a list of indices: arr[[0, 2]] or arr[0, 2]
# Try casting to 1D integer tensor if all elements are ints
try:
key_list = [torch.tensor(key, device=device)]
except (RuntimeError, TypeError, ValueError):
key_list = list(key)
else:
key_list = [key]
normalized = []
for k in key_list:
# Convert numpy array to torch tensor on target device
if isinstance(k, np.ndarray):
normalized.append(torch.from_numpy(k).to(device=device))
# Unwrap 0-D scalar DNDarray into a Python scalar
elif isinstance(k, DNDarray) and k.ndim == 0:
normalized.append(k.larray.item())
# Unpack singleton containers like (idx,) often produced by nonzero/where
elif isinstance(k, (tuple, list)) and len(k) == 1 and isinstance(k[0], DNDarray):
if k[0].ndim > 0:
normalized.append(k[0])
else:
normalized.append(torch.tensor([k[0].larray.item()], device=device))
# Sequence of scalar DNDarrays -> unwrap to list of Python scalars
elif (
isinstance(k, (tuple, list))
and len(k) > 0
and all(isinstance(elem, DNDarray) and elem.ndim == 0 for elem in k)
):
normalized.append(torch.tensor([elem.larray.item() for elem in k], device=device))
# Catch invalid nested non-scalar DNDarrays early
elif isinstance(k, (tuple, list)) and any(
isinstance(elem, DNDarray) and elem.ndim > 0 for elem in k
):
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])."
)
# Convert non-distributed integer/indexing DNDarrays to local torch.Tensor
elif isinstance(k, DNDarray) and k.split is None and k.dtype in (types.int32, types.int64):
normalized.append(k.larray.to(dtype=torch.int64))
# Ensure torch.Tensor indices are placed on the target device
elif isinstance(k, torch.Tensor):
normalized.append(k)
else:
# Leave slices, integers, None, Ellipsis,
# and distributed DNDarrays (ndim >= 1 and split is not None) intact.
normalized.append(k)
return tuple(normalized)
def _scalar_early_out(
arr: "DNDarray",
key: Any,
op: str | None,
return_local_indices: bool | None,
) -> tuple["DNDarray", ProcessedKey]:
"""Resolve early-out for scalar indexers."""
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
processed_key, root = _process_scalar_key(
arr, key, indexed_axis=0, return_local_indices=return_local_indices
)
return arr, ProcessedKey(
key=processed_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 if output_split is None else arr.balanced,
root=root,
)
def _expand_dimensions_and_ellipsis(
arr: "DNDarray",
key: list[Any],
output_shape: list[int],
split_bookkeeping: list[str | None],
) -> tuple["DNDarray", list[Any], list[int], list[str | None]]:
"""
Expands ellipses (...) into full slices and inserts singleton dimensions
for None (newaxis) or 0-D boolean masks.
"""
add_dims = sum(k is None or _is_boolean_scalar(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:
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:
for i, k in reversed(list(enumerate(key))):
if k is None or _is_boolean_scalar(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)
insert_pos = i - add_dims + 1
arr = arr.expand_dims(insert_pos)
output_shape = output_shape[:insert_pos] + [1] + output_shape[insert_pos:]
split_bookkeeping = (
split_bookkeeping[:insert_pos] + [None] + split_bookkeeping[insert_pos:]
)
add_dims -= 1
return arr, key, output_shape, split_bookkeeping
def _distr_mask_fast_path(arr: "DNDarray", key: Any, op: str | None) -> bool:
"""
Checks if the indexing operation qualifies for the distributed boolean mask fast path.
"""
if not arr.is_distributed():
return False
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
):
if split_key.gshape == arr.gshape:
# "get" flattens to 1D; if split > 0, local flattening scrambles global C-order
return op == "set" or (op == "get" and arr.split == 0)
elif (
split_key.ndim == 1 and arr.split == 0 and split_key.gshape == (arr.gshape[arr.split],)
):
return True
return False
def _resolve_1d_boolean_first_dim(
arr: "DNDarray", key: tuple[Any, ...] | Any, distr_mask_fast_path: bool
) -> tuple[Any, ...] | Any:
"""
If key indexes axis 0 with a 1D boolean mask matching the global size of axis 0,
convert that mask into integer coordinates via nonzero().
"""
if distr_mask_fast_path or arr.ndim == 0:
return key
first = key[0] if isinstance(key, tuple) and len(key) >= 1 else key
if not isinstance(first, (DNDarray, torch.Tensor)):
return key
first_dtype = getattr(first, "dtype", None)
first_ndim = getattr(first, "ndim", 0)
first_shape = tuple(getattr(first, "shape", ()))
if (
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()
idx0 = nz[0] if isinstance(nz, tuple) else nz
elif isinstance(first, torch.Tensor):
idx0 = torch.nonzero(first, as_tuple=False).flatten()
return (idx0,) + key[1:] if isinstance(key, tuple) else (idx0,)
return key
def _sanitize_int_indices(k: "DNDarray", dim: int, axis: int, comm: Any, device: Any) -> "DNDarray":
"""
Validates integer bounds and normalizes negative coordinates for distributed/local DNDarray keys.
"""
if k.dtype not in (types.int32, types.int64):
if k.dtype not in (ht_bool, ht_uint8):
raise IndexError(
f"arrays used as indices must be of integer (or boolean) type, got {k.dtype}"
)
return k
# Combine local checks into one reduced boolean tensor
local_flags = torch.tensor(
[
((k.larray < -dim) | (k.larray >= dim)).any(),
(k.larray < 0).any(),
],
dtype=torch.int32,
device=device.torch_device,
)
do_reduce = comm is not None and getattr(comm, "size", 1) > 1 and k.is_distributed()
if do_reduce:
comm.Allreduce(MPI.IN_PLACE, local_flags, op=MPI.SUM)
invalid_sum = local_flags[0].item()
has_neg_sum = local_flags[1].item()
if invalid_sum > 0:
raise IndexError(f"index out of bounds for axis {axis} 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=device,
comm=comm,
copy=False,
)
return k
def _process_slice_indexer(
k: slice,
dim: int,
is_split_axis: bool,
displs: list[int] | None,
counts: list[int] | None,
rank: int | None,
device: torch.device,
return_local_indices: bool,
) -> tuple[Any, int, int | None, bool | None]:
"""
Computes local slice or index tensor along an axis, determining output dimension length
and slice monotonicity ordering.
"""
if k.step == 0:
raise ValueError("Slice step cannot be zero")
start, stop, step = slice(k.start, k.stop, k.step).indices(dim)
new_key = k
output_dim_len = 0
split_key_is_ordered = None
out_is_balanced = None
if step < 0 and start > stop:
# total items in the global descending slice
output_dim_len = len(range(start, stop, step))
if is_split_axis:
split_key_is_ordered = -1
out_is_balanced = False
# PyTorch cannot index with negative steps
# we work with the values in ascending order
s = -step
# lowest global coordinate produced by this slice
min_coord = start - (output_dim_len - 1) * s
# highest global coordinate produced by this slice
max_coord = start
# global index interval [low, high) owned by this MPI process
low = displs[rank]
high = low + counts[rank]
# overlap between the slice value bounds [min_coord, max_coord + 1)
# and the current process's memory chunk [low, high)
overlap_start = max(min_coord, low)
overlap_end = min(max_coord + 1, high)
# check if there is any overlap at all
if overlap_start < overlap_end:
# first item in the progression that is >= overlap_start.
k_first = (overlap_start - min_coord + s - 1) // s
g_first = min_coord + k_first * s
if g_first < overlap_end:
# first value is still inside the rank's chunk
# calculate how many steps fit in [g_first, overlap_end)
steps_local = (overlap_end - 1 - g_first) // s + 1
# convert global index to a process-local offset if requested
start_idx = g_first - low if return_local_indices else g_first
stop_idx = start_idx + steps_local * s
# allocate local indices only
new_key = slice(start_idx, stop_idx, s)
else:
# slicing skips local chunk completely
new_key = slice(0, 0)
else:
# local chunk is outside the slice bounds entirely
new_key = slice(0, 0)
else:
# non-split axis: return the descending slice as indices
new_key = torch.arange(start, stop, step, device=device, dtype=torch.int64)
elif step > 0 and start < stop:
output_dim_len = len(range(start, stop, step))
if is_split_axis:
split_key_is_ordered = 1
out_is_balanced = False
local_arr_end = displs[rank] + counts[rank]
if stop > displs[rank] and start < local_arr_end:
index_in_cycle = (displs[rank] - start) % step
if start >= displs[rank]:
local_start = start - displs[rank]
else:
local_start = 0 if index_in_cycle == 0 else step - index_in_cycle
if stop <= local_arr_end:
local_stop = stop - displs[rank]
else:
local_stop = counts[rank]
new_key = slice(local_start, local_stop, step)
else:
new_key = slice(0, 0)
else:
new_key = slice(0, 0)
output_dim_len = 0
return new_key, output_dim_len, split_key_is_ordered, out_is_balanced
def _reorder_advanced_idx_axes(
arr: "DNDarray",
key: list[Any],
advanced_indexing_dims: list[int],
advanced_indexing_shapes: list[tuple[int, ...]],
output_shape: list[int | None],
split_bookkeeping: list[str | None],
key_is_mask_like: bool,
) -> tuple["DNDarray", list[Any], list[int | None], list[str | None], tuple[int, ...]]:
"""
Broadcasts advanced indexing dimensions and rearranges non-consecutive dimensions
to the front of the array as mandated by NumPy advanced indexing semantics.
"""
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)
is_consecutive = (
len(advanced_indexing_dims) == 1
or list(range(advanced_indexing_dims[0], advanced_indexing_dims[-1] + 1))
== advanced_indexing_dims
)
if is_consecutive:
output_shape[
advanced_indexing_dims[0] : advanced_indexing_dims[0] + len(advanced_indexing_dims)
] = broadcasted_shape
if key_is_mask_like:
has_split = (
"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"] if has_split else [None]
else:
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:
new_idx = max(0, adv_sb.index("split") + add_dims)
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:
# Non-consecutive: transpose to make advanced dims leading and consecutive
non_adv_ind_dims = [i for i in range(arr.ndim) if i not in advanced_indexing_dims]
transpose_axes = tuple(advanced_indexing_dims + non_adv_ind_dims)
arr = arr.transpose(transpose_axes)
output_shape = [output_shape[i] for i in transpose_axes]
output_shape[: len(advanced_indexing_dims)] = broadcasted_shape
split_bookkeeping = [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 = max(0, adv_sb.index("split") + add_dims)
new_adv_sb[new_idx] = "split"
split_bookkeeping = new_adv_sb + split_bookkeeping[len(advanced_indexing_dims) :]
key = [key[i] for i in advanced_indexing_dims] + [key[i] for i in non_adv_ind_dims]
return arr, key, output_shape, split_bookkeeping
def _assess_op_type(
root: int | None,
split_key_is_ordered: int,
distr_mask_fast_path: bool,
key_is_mask_like: bool,
) -> str:
"""Determine the indexing operation routing category."""
if root is not None:
return "scalar"
if split_key_is_ordered == 0:
return "distributed"
if split_key_is_ordered == -1:
return "descending_slice"
if distr_mask_fast_path:
return "distr_mask"
if key_is_mask_like:
return "local_mask"
return "local"
def _sanitize_advanced_keys(
arr: "DNDarray",
key: list[Any],
advanced_indexing_dims: list[int],
split_key_is_ordered: int,
key_is_mask_like: bool,
distr_mask_fast_path: bool,
counts: tuple | list | None,
displs: tuple | list | None,
return_local_indices: bool,
) -> tuple[list[Any], bool]:
"""
Validates key distribution alignment along the split axis and converts
all advanced indexing key elements from DNDarrays into local torch.Tensors.
"""
key = list(key)
# Detect mask-like conditions (same shape for adv indexing dimensions)
adv_keys = [key[i] for i in advanced_indexing_dims]
key_is_mask_like = key_is_mask_like or (
len(advanced_indexing_dims) > 1
and all(isinstance(k, DNDarray) for k in adv_keys)
and len(set(k.shape for k in adv_keys)) == 1
)
non_split_dims = [d for d in advanced_indexing_dims if d != arr.split]
# Align distributions if mask-like
if key_is_mask_like and arr.split is not None and arr.split in advanced_indexing_dims:
key_splits = [k.split for k in adv_keys]
split_pos = advanced_indexing_dims.index(arr.split)
target_split = key_splits[split_pos]
if key_splits.count(target_split) != len(key_splits):
if target_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=target_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}."
)
# Extract local torch.Tensors
if arr.is_distributed() and arr.split in advanced_indexing_dims:
if distr_mask_fast_path:
for i in non_split_dims:
if isinstance(key[i], DNDarray):
key[i] = key[i].larray
elif split_key_is_ordered == 1:
k = key[arr.split].larray if isinstance(key[arr.split], DNDarray) else key[arr.split]
rank = arr.comm.rank
low = displs[rank]
high = low + counts[rank]
idx_start = torch.searchsorted(k, low)
idx_end = torch.searchsorted(k, high)
k_local = k[idx_start:idx_end]
if return_local_indices:
k_local = k_local - low
key[arr.split] = k_local
if key_is_mask_like:
for i in non_split_dims:
larr = key[i].larray if isinstance(key[i], DNDarray) else key[i]
key[i] = larr[idx_start:idx_end]
else:
for i in non_split_dims:
if isinstance(key[i], DNDarray):
key[i] = key[i].larray
else:
# split_key_is_ordered == 0 (unordered indexing)
for i in advanced_indexing_dims:
if isinstance(key[i], DNDarray):
key[i] = key[i].larray
else:
for i in advanced_indexing_dims:
if isinstance(key[i], DNDarray):
key[i] = key[i].larray
return key, key_is_mask_like
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"``, ``"local"``, 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.
"""
# early out for scalar key
if _is_scalar_index(key):
return _scalar_early_out(
arr=arr,
key=key,
op=op,
return_local_indices=return_local_indices,
)
# normalize key items to torch-friendly types (torch.Tensor, int, slice, None, Ellipsis)
# NB: distributed DNDarrays are not unwrapped here, they are handled later
normalized_key = _normalize_key(key, device=arr.device.torch_device)
# maintain single item when raw key was not passed as a tuple
key = normalized_key if isinstance(key, tuple) else normalized_key[0]
# unpack single-element tuple containing a boolean mask (a[(mask,)] same as a[mask])
if isinstance(key, tuple) and len(key) == 1 and _is_boolean_array(key[0]):
key = key[0]
# evaluate if this is a distributed mask aligned with the array
distr_mask_fast_path = _distr_mask_fast_path(arr, key, op)
if distr_mask_fast_path and not isinstance(key, tuple):
return arr, ProcessedKey(
key=key.larray,
op_type="distr_mask",
output_shape=None,
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,
)
# 1D boolean mask resolution
key = _resolve_1d_boolean_first_dim(arr, key, distr_mask_fast_path)
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
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:
split_key_is_ordered = int((key[1:] >= key[:-1]).all().item())
# unordered local keys
if not split_key_is_ordered and not key_is_dist:
out_is_balanced = True
# ordered keys
if split_key_is_ordered:
rank = arr.comm.rank
low = displs[rank]
high = low + counts[rank]
idx_start = torch.searchsorted(key, low)
idx_end = torch.searchsorted(key, high)
key = key[idx_start:idx_end]
if return_local_indices:
key = key - low
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
op_type = _assess_op_type(
root=root,
split_key_is_ordered=split_key_is_ordered,
distr_mask_fast_path=distr_mask_fast_path,
key_is_mask_like=key_is_mask_like,
)
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,
)
if isinstance(key, (tuple, list)):
key = list(key)
else:
key = [key]
# check for ellipsis, newaxis. NB: (ht.newaxis is None)==True
arr, key, output_shape, split_bookkeeping = _expand_dimensions_and_ellipsis(
arr, key, output_shape, split_bookkeeping
)
# recalculate new_split, transpose_axes after dimensions manipulation
new_split = split_bookkeeping.index("split") if "split" in split_bookkeeping else None
# check for advanced indexing and slices
advanced_indexing_dims = []
advanced_indexing_shapes = []
for i, k in enumerate(key):
if _is_scalar_index(k):
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"
)
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
k = _sanitize_int_indices(
k=k, dim=arr.gshape[i], axis=i, comm=arr.comm, device=arr.device
)
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.numel() <= 1 or (k.larray[1:] >= k.larray[:-1]).all().item())
):
split_key_is_ordered = 1
out_is_balanced = False
else:
split_key_is_ordered = 0
out_is_balanced = True
key[i] = k
elif isinstance(k, slice) and k != slice(None):
is_split_axis = arr_is_distributed and new_split == i
rank = arr.comm.rank if arr_is_distributed else None
slice_key, dim_len, s_ordered, s_balanced = _process_slice_indexer(
k=k,
dim=arr.gshape[i],
is_split_axis=is_split_axis,
displs=displs if arr_is_distributed else None,
counts=counts if arr_is_distributed else None,
rank=rank,
device=arr.device.torch_device,
return_local_indices=return_local_indices,
)
key[i] = slice_key
output_shape[i] = dim_len
if is_split_axis and s_ordered is not None:
split_key_is_ordered = s_ordered
out_is_balanced = s_balanced
if advanced_indexing:
key, key_is_mask_like = _sanitize_advanced_keys(
arr=arr,
key=key,
advanced_indexing_dims=advanced_indexing_dims,
split_key_is_ordered=split_key_is_ordered,
key_is_mask_like=key_is_mask_like,
distr_mask_fast_path=distr_mask_fast_path,
counts=counts if arr_is_distributed else None,
displs=displs if arr_is_distributed else None,
return_local_indices=return_local_indices,
)
arr, key, output_shape, split_bookkeeping = _reorder_advanced_idx_axes(
arr=arr,
key=key,
advanced_indexing_dims=advanced_indexing_dims,
advanced_indexing_shapes=advanced_indexing_shapes,
output_shape=output_shape,
split_bookkeeping=split_bookkeeping,
key_is_mask_like=key_is_mask_like,
)
# expand key to match the number of dimensions of the DNDarray
if arr.ndim > len(key):
key += [slice(None)] * (arr.ndim - len(key))
while None in output_shape:
lost_dim = output_shape.index(None)
output_shape.pop(lost_dim)
split_bookkeeping.pop(lost_dim)
output_shape = tuple(output_shape)
new_split = split_bookkeeping.index("split") if "split" in split_bookkeeping else None
op_type = _assess_op_type(
root=root,
split_key_is_ordered=split_key_is_ordered,
distr_mask_fast_path=distr_mask_fast_path,
key_is_mask_like=key_is_mask_like,
)
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,
)
[docs]
class DNDarray:
"""
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,...]
The global shape of the array
dtype : datatype
The datatype of the array
split : int or None
The axis on which the array is divided between processes
device : Device
The device on which the local arrays are using (cpu or gpu)
comm : Communication
The communications object for sending and receiving data
balanced: bool or None
Describes whether the data are evenly distributed across processes.
If this information is not available (``self.balanced is None``), it
can be gathered via the :func:`is_balanced()` method (requires communication).
"""
def __init__(
self,
array: torch.Tensor,
gshape: tuple[int, ...],
dtype: datatype,
split: int | None,
device: Device,
comm: Communication,
balanced: bool,
):
self.__array = array
self.__gshape = gshape
self.__dtype = dtype
self.__split = split
self.__device = device
self.__comm = comm
self.__balanced: bool = balanced
self.__ishalo = False
self.__halo_next: torch.Tensor | None = None
self.__halo_prev: torch.Tensor | None = None
self.__partitions_dict__ = None
self.__lshape_map = None
self.__counts_displs = None
# check for inconsistencies between torch and heat devices
assert str(array.device) == device.torch_device
@property
def balanced(self) -> bool:
"""
Boolean value indicating if the DNDarray is balanced between the MPI processes
"""
return self.__balanced
@property
def comm(self) -> Communication:
"""
The :class:`~heat.core.communication.Communication` of the ``DNDarray``
"""
return self.__comm
@property
def device(self) -> Device:
"""
The :class:`~heat.core.devices.Device` of the ``DNDarray``
"""
return self.__device
@property
def dtype(self) -> datatype:
"""
The :class:`~heat.core.types.datatype` of the ``DNDarray``
"""
return self.__dtype
@property
def gshape(self) -> tuple:
"""
Returns the global shape of the ``DNDarray`` across all processes
"""
return self.__gshape
@property
def halo_next(self) -> torch.Tensor:
"""
Returns the halo of the next process
"""
return self.__halo_next
@property
def halo_prev(self) -> torch.Tensor:
"""
Returns the halo of the previous process
"""
return self.__halo_prev
@property
def larray(self) -> torch.Tensor:
"""
Returns the underlying process-local ``torch.Tensor`` of the ``DNDarray``
"""
return self.__array
@larray.setter
def larray(self, array: torch.Tensor):
"""
Setter for ``self.larray``, the underlying local ``torch.Tensor`` of the ``DNDarray``.
Parameters
----------
array : torch.Tensor
The new underlying local ``torch.tensor`` of the ``DNDarray``
Warning
-----------
Please use this function with care, as it might corrupt/invalidate the metadata in the ``DNDarray`` instance.
"""
# sanitize tensor input
sanitation.sanitize_in_tensor(array)
# verify consistency of tensor shape with global DNDarray
sanitation.sanitize_lshape(self, array)
# set balanced status
split = self.split
if split is not None and array.shape[split] != self.lshape[split]:
self.__balanced = None
self.__lshape_map = None
self.__counts_displs = None
self.__array = array
@property
def nbytes(self) -> int:
"""
Returns the number of bytes consumed by the global tensor. Equivalent to property gnbytes.
Note
------------
Does not include memory consumed by non-element attributes of the ``DNDarray`` object.
"""
return self.__array.element_size() * self.size
@property
def ndim(self) -> int:
"""
Number of dimensions of the ``DNDarray``
"""
return len(self.__gshape)
@property
def __partitioned__(self) -> dict:
"""
Return a dictionary containing information useful for working with the partitioned
data. These items include the shape of the data on each process, the starting index of the data
that a process has, the datatype of the data, the local devices, as well as the global
partitioning scheme.
An example of the output and shape is shown in :func:`ht.core.DNDarray.create_partition_interface <ht.core.DNDarray.create_partition_interface>`.
Returns
-------
dictionary with the partition interface
"""
if self.__partitions_dict__ is None:
self.__partitions_dict__ = self.create_partition_interface()
return self.__partitions_dict__
@property
def size(self) -> int:
"""
Number of total elements of the ``DNDarray``
"""
if self.larray.is_mps:
# MPS does not support double precision
size = torch.prod(
torch.tensor(self.gshape, dtype=torch.float32, device=self.device.torch_device)
)
else:
size = torch.prod(
torch.tensor(self.gshape, dtype=torch.float64, device=self.device.torch_device)
)
return size.long().item()
@property
def gnbytes(self) -> int:
"""
Returns the number of bytes consumed by the global ``DNDarray``
Note
-----------
Does not include memory consumed by non-element attributes of the ``DNDarray`` object.
"""
return self.nbytes
@property
def gnumel(self) -> int:
"""
Returns the number of total elements of the ``DNDarray``
"""
return self.size
@property
def imag(self) -> DNDarray:
"""
Return the imaginary part of the ``DNDarray``.
"""
return complex_math.imag(self)
@property
def lnbytes(self) -> int:
"""
Returns the number of bytes consumed by the local ``torch.Tensor``
Note
-------------------
Does not include memory consumed by non-element attributes of the ``DNDarray`` object.
"""
return self.__array.element_size() * self.__array.nelement()
@property
def lnumel(self) -> int:
"""
Number of elements of the ``DNDarray`` on each process
"""
return np.prod(self.__array.shape)
@property
def lshape(self) -> tuple[int]:
"""
Returns the shape of the ``DNDarray`` on each node
"""
return tuple(self.__array.shape)
@property
def lshape_map(self) -> torch.Tensor:
"""
Returns the lshape map. If it hasn't been previously created then it will be created here.
"""
return self.create_lshape_map()
@property
def lloc(self):
"""Deprecated function for local indexing. Use `DNDarray.larray` for local indexing instead"""
# TODO: Remove this entirely by heat v2.5
raise Exception(
"`DNDarray.lloc` is deprecated. Use `DNDarray.larray` for local indexing instead."
)
@property
def real(self) -> DNDarray:
"""
Return the real part of the ``DNDarray``.
"""
return complex_math.real(self)
@property
def shape(self) -> tuple[int]:
"""
Returns the shape of the ``DNDarray`` as a whole
"""
return self.__gshape
@property
def split(self) -> int | None:
"""
Returns the axis on which the ``DNDarray`` is split
"""
return self.__split
@property
def stride(self) -> tuple[int]:
"""
Returns the steps in each dimension when traversing a ``DNDarray``. torch-like usage: ``self.stride()``
"""
return self.__array.stride
@property
def strides(self) -> tuple[int]:
"""
Returns bytes to step in each dimension when traversing a ``DNDarray``. numpy-like usage: ``self.strides()``
"""
steps = list(self.__array.stride())
try:
itemsize = self.__array.untyped_storage().element_size()
except AttributeError:
itemsize = self.__array.storage().element_size()
strides = tuple(step * itemsize for step in steps)
return strides
@property
def array_with_halos(self) -> torch.Tensor:
"""
Fetch halos of size ``halo_size`` from neighboring ranks and save them in ``self.halo_next``/``self.halo_prev``
in case they are not already stored. If ``halo_size`` differs from the size of already stored halos,
the are overwritten.
"""
return self.__cat_halo()
def __prephalo(self, start, end) -> torch.Tensor:
"""
Extracts the halo indexed by start, end from ``self.array`` in the direction of ``self.split``
Parameters
----------
start : int
Start index of the halo extracted from ``self.array``
end : int
End index of the halo extracted from ``self.array``
"""
ix = [slice(None, None, None)] * len(self.shape)
try:
ix[self.split] = slice(start, end)
except IndexError:
print("Indices out of bound")
return self.__array[tuple(ix)].clone()
[docs]
def get_halo(self, halo_size: int, prev: bool = True, next: bool = True):
"""
Fetch halos of size ``halo_size`` from neighboring ranks and save them in ``self.halo_next/self.halo_prev``.
Parameters
----------
halo_size : int
Size of the halo.
prev : bool, optional
If True, fetch the halo from the previous rank. Default: True.
next : bool, optional
If True, fetch the halo from the next rank. Default: True.
"""
if not isinstance(halo_size, int):
raise TypeError(
f"halo_size needs to be of Python type integer, {type(halo_size)} given"
)
if halo_size < 0:
raise ValueError(
f"halo_size needs to be a non-negative Python integer, {halo_size} given"
)
if self.is_distributed() and halo_size > 0:
# gather lshapes
lshape_map = self.lshape_map
rank = self.comm.rank
populated_ranks = torch.nonzero(lshape_map[:, self.split]).squeeze().tolist()
if rank in populated_ranks:
first_rank = populated_ranks[0]
last_rank = populated_ranks[-1]
if rank != last_rank:
next_rank = populated_ranks[populated_ranks.index(rank) + 1]
if rank != first_rank:
prev_rank = populated_ranks[populated_ranks.index(rank) - 1]
else:
# if process has no data we ignore it
return
if (halo_size > self.lshape_map[:, self.split][populated_ranks]).any():
# halo_size is larger than the local size on at least one process
raise ValueError(
f"halo_size {halo_size} needs to be smaller than chunk-size {self.lshape[self.split]} )"
)
a_prev = self.__prephalo(0, halo_size)
a_next = self.__prephalo(-halo_size, None)
res_prev = None
res_next = None
req_list = []
# exchange data with next populated process
if prev:
if rank != last_rank:
req_list.append(self.comm.Isend(a_next, next_rank))
if rank != first_rank:
res_prev = torch.empty(
a_prev.size(), dtype=a_prev.dtype, device=self.device.torch_device
)
req_list.append(self.comm.Irecv(res_prev, source=prev_rank))
if next:
if rank != first_rank:
req_list.append(self.comm.Isend(a_prev, prev_rank))
if rank != last_rank:
res_next = torch.empty(
a_next.size(), dtype=a_next.dtype, device=self.device.torch_device
)
req_list.append(self.comm.Irecv(res_next, source=next_rank))
for req in req_list:
req.Wait()
self.__halo_next = res_next
self.__halo_prev = res_prev
self.__ishalo = True
def __cat_halo(self) -> torch.Tensor:
"""
Return local array concatenated to halos if they are available.
"""
if not self.is_distributed():
return self.__array
return torch.cat(
[_ for _ in (self.__halo_prev, self.__array, self.__halo_next) if _ is not None],
dim=self.split,
)
[docs]
def __array__(self) -> np.ndarray:
"""
Returns a view of the process-local slice of the :class:`DNDarray` as a numpy ndarray, if the ``DNDarray`` resides on CPU. Otherwise, it returns a copy, on CPU, of the process-local slice of ``DNDarray`` as numpy ndarray.
"""
return self.larray.cpu().__array__()
[docs]
def __array_ufunc__(self, ufunc, method, *inputs, **kwargs):
"""
Override NumPy's universal functions.
"""
import heat
# TODO support ufunc method variants
if method == "__call__":
try:
func = getattr(heat, ufunc.__name__)
except AttributeError:
return NotImplemented
return func(*inputs, **kwargs)
else:
return NotImplemented
[docs]
def __array_function__(self, func, types, args, kwargs):
"""
Augments NumPy's functions.
"""
import heat
try:
ht_func = getattr(heat, func.__name__)
except AttributeError:
return NotImplemented
return ht_func(*args, **kwargs)
[docs]
def __array_namespace__(self, *, api_version: str | None = None) -> Any:
"""
Returns an object that has all the array API functions on it.
Parameters
----------
api_version : Optional[str]
string representing the version of the array API specification to
be returned, in ``'YYYY.MM'`` form. If it is ``None`` (default), it
returns the namespace corresponding to latest version of the
array API specification.
"""
if api_version is not None and api_version != "2025.12":
raise ValueError(f"Unrecognized array API version: {api_version}")
import heat
return heat
[docs]
def astype(self, dtype, copy=True, device: Device = None) -> DNDarray:
"""
Returns a casted version of this array.
Casted array is a new array of the same shape but with given type of this array. If copy is ``True``, the
same array is returned instead.
Parameters
----------
dtype : datatype
Heat type to which the array is cast
copy : bool, optional
By default the operation returns a copy of this array. If copy is set to ``False`` the cast is performed
in-place and this array is returned
device: ht.Device, optional
The device on which to place the array. If ``None``, keep device. Default: None.
"""
dtype = canonical_heat_type(dtype)
device = self.__device if device is None else devices.sanitize_device(device)
if self.__array.is_mps:
if dtype == types.float64:
# print warning
warnings.warn(
"MPS does not support float64. Casting to float32 instead.",
ResourceWarning,
)
dtype = types.float32
elif dtype == types.complex128:
# print warning
warnings.warn(
"MPS does not support complex128. Casting to complex64 instead.",
ResourceWarning,
)
dtype = types.complex64
casted_array = self.__array.to(
device=device.torch_device, dtype=dtype.torch_type(), copy=copy
)
if copy:
return DNDarray(
casted_array,
gshape=self.shape,
dtype=dtype,
split=self.split,
device=device,
comm=self.comm,
balanced=self.balanced,
)
self.__array = casted_array
self.__dtype = dtype
self.__device = device
return self
[docs]
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``
itself and will not return anything.
Examples
--------
>>> a = ht.zeros((10, 2), split=0)
>>> a[:, 0] = ht.arange(10)
>>> b = a[3:]
[0/2] tensor([[3., 0.],
[1/2] tensor([[4., 0.],
[5., 0.],
[6., 0.]])
[2/2] tensor([[7., 0.],
[8., 0.],
[9., 0.]])
>>> b.balance_()
>>> print(b.gshape, b.lshape)
[0/2] (7, 2) (1, 2)
[1/2] (7, 2) (3, 2)
[2/2] (7, 2) (3, 2)
>>> b
[0/2] tensor([[3., 0.],
[4., 0.],
[5., 0.]])
[1/2] tensor([[6., 0.],
[7., 0.]])
[2/2] tensor([[8., 0.],
[9., 0.]])
>>> print(b.gshape, b.lshape)
[0/2] (7, 2) (3, 2)
[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_()
[docs]
def __bool__(self) -> bool:
"""
Boolean scalar casting.
"""
return self.__cast(bool)
def __cast(self, cast_function) -> float | int:
"""
Implements a generic cast function for ``DNDarray`` objects.
Parameters
----------
cast_function : function
The actual cast function, e.g. ``float`` or ``int``
Raises
------
TypeError
If the ``DNDarray`` object cannot be converted into a scalar.
"""
if np.prod(self.shape) == 1:
if not self.is_distributed():
return cast_function(self.__array)
is_empty = np.prod(self.__array.shape) == 0
root = self.comm.allreduce(0 if is_empty else self.comm.rank, op=MPI.SUM)
return self.comm.bcast(None if is_empty else cast_function(self.__array), root=root)
raise TypeError("only size-1 arrays can be converted to Python scalars")
[docs]
def collect_(self, target_rank: int = 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.
Parameters
----------
target_rank : int, optional
The rank to which the DNDarray will be collected. Default: 0.
Raises
------
TypeError
If the target rank is not an integer.
ValueError
If the target rank is out of bounds.
Examples
--------
>>> st = ht.ones((50, 81, 67), split=2)
>>> print(st.lshape)
[0/2] (50, 81, 23)
[1/2] (50, 81, 22)
[2/2] (50, 81, 22)
>>> st.collect_()
>>> print(st.lshape)
[0/2] (50, 81, 67)
[1/2] (50, 81, 0)
[2/2] (50, 81, 0)
>>> st.collect_(1)
>>> print(st.lshape)
[0/2] (50, 81, 0)
[1/2] (50, 81, 67)
[2/2] (50, 81, 0)
"""
if not isinstance(target_rank, int):
raise TypeError(f"target rank must be of type int , but was {type(target_rank)}")
if target_rank >= self.comm.size:
raise ValueError("target rank is out of bounds")
if not self.is_distributed():
return
target_map = self.lshape_map.clone()
target_map[:, self.split] = 0
target_map[target_rank, self.split] = self.gshape[self.split]
self.redistribute_(target_map=target_map)
[docs]
def __complex__(self) -> DNDarray:
"""
Complex scalar casting.
"""
return self.__cast(complex)
[docs]
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.
"""
if self.split is not None:
if self.__counts_displs is not None:
return self.__counts_displs
if self.__lshape_map is None:
self.create_lshape_map()
counts = self.__lshape_map[:, self.split]
displs = [0] + torch.cumsum(counts, dim=0)[:-1].tolist()
res = (tuple(counts.tolist()), tuple(displs))
self.__counts_displs = res
return res
raise ValueError("Non-distributed DNDarray. Cannot calculate counts and displacements.")
[docs]
def cpu(self) -> DNDarray:
"""
Returns a copy of this object in main memory. If this object is already in main memory, then no copy is
performed and the original object is returned.
"""
self.__array = self.__array.cpu()
self.__device = devices.cpu
return self
[docs]
def create_lshape_map(self, force_check: bool = False) -> torch.Tensor:
"""
Generate a 'map' of the lshapes of the data on all processes.
Units are ``(process rank, lshape)``
Parameters
----------
force_check : bool, optional
if False (default) and the lshape map has already been created, use the previous
result. Otherwise, create the lshape_map
"""
if not force_check and self.__lshape_map is not None:
return self.__lshape_map.clone()
lshape_map = torch.zeros(
(self.comm.size, self.ndim), dtype=torch.int64, device=self.device.torch_device
)
if not self.is_distributed():
lshape_map[:] = torch.tensor(self.gshape, device=self.device.torch_device)
self.__lshape_map = lshape_map
return lshape_map.clone()
elif self.is_balanced(force_check=True):
for i in range(self.comm.size):
_, lshape, _ = self.comm.chunk(self.gshape, self.split, rank=i)
lshape_map[i, :] = torch.tensor(lshape, device=self.device.torch_device)
else:
lshape_map[self.comm.rank, :] = torch.tensor(
self.lshape, device=self.device.torch_device
)
self.comm.Allreduce(MPI.IN_PLACE, lshape_map, MPI.SUM)
self.__lshape_map = lshape_map
self.__counts_displs = None
return lshape_map.clone()
[docs]
def create_partition_interface(self):
"""
Create a partition interface in line with the DPPY proposal. This is subject to change.
The intention of this to facilitate the usage of a general format for the referencing of
distributed datasets.
An example of the output and shape is shown below.
__partitioned__ = {
'shape': (27, 3, 2),
'partition_tiling': (4, 1, 1),
'partitions': {
(0, 0, 0): {
'start': (0, 0, 0),
'shape': (7, 3, 2),
'data': tensor([...], dtype=torch.int32),
'location': [0],
'dtype': torch.int32,
'device': 'cpu'
},
(1, 0, 0): {
'start': (7, 0, 0),
'shape': (7, 3, 2),
'data': None,
'location': [1],
'dtype': torch.int32,
'device': 'cpu'
},
(2, 0, 0): {
'start': (14, 0, 0),
'shape': (7, 3, 2),
'data': None,
'location': [2],
'dtype': torch.int32,
'device': 'cpu'
},
(3, 0, 0): {
'start': (21, 0, 0),
'shape': (6, 3, 2),
'data': None,
'location': [3],
'dtype': torch.int32,
'device': 'cpu'
}
},
'locals': [(rank, 0, 0)],
'get': lambda x: x,
}
Returns
-------
dictionary containing the partition interface as shown above.
"""
lshape_map = self.create_lshape_map()
start_idx_map = torch.zeros_like(lshape_map)
part_tiling = [1] * self.ndim
lcls = [0] * self.ndim
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:
start_idx_map[:] = 0
partitions = {}
base_key = [0] * self.ndim
for r in range(self.comm.size):
if self.split is not None:
base_key[self.split] = r
dat = None if r != self.comm.rank else self.larray
else:
dat = self.larray
partitions[tuple(base_key)] = {
"start": tuple(start_idx_map[r].tolist()),
"shape": tuple(lshape_map[r].tolist()),
"data": dat,
"location": [r],
"dtype": self.dtype.torch_type(),
"device": self.device.torch_device,
}
partition_dict = {
"shape": self.gshape,
"partition_tiling": tuple(part_tiling),
"partitions": partitions,
"locals": [tuple(lcls)],
"get": lambda x: x,
}
self.__partitions_dict__ = partition_dict
return partition_dict
[docs]
def __dlpack__(
self,
*args,
**kwargs,
) -> Any:
"""
Exports the undistributed array for consumption by ``from_dlpack()`` as a DLPack capsule.
Any positional arguments ``*args`` and keyword arguments ``**kwargs`` are directly forwarded to torch ``__dlpack__``.
Note
----
See `Array API <https://data-apis.org/array-api/2025.12/API_specification/generated/array_api.array.__dlpack__.html>`_ for details and the function signature as implemented by torch.
Raises
------
BufferError
if the DNDarray is distributed, as this is not supported by DLPack.
"""
if self.is_distributed():
raise BufferError("DLPack export works for undistributed arrays only.")
return self.larray.__dlpack__(*args, **kwargs)
[docs]
def __dlpack_device__(self) -> tuple[Enum, int]:
"""
Returns device type and device ID in DLPack format. Meant for use
within ``from_dlpack()``.
"""
return self.larray.__dlpack_device__()
[docs]
def __float__(self) -> float:
"""
Float scalar casting.
See Also
--------
:func:`~heat.core.manipulations.flatten`
"""
return self.__cast(float)
[docs]
def fill_diagonal(self, value: float) -> DNDarray:
"""
Fill the main diagonal of a 2D :class:`DNDarray`.
This function modifies the input tensor in-place, and returns the input array.
Parameters
----------
value : float
The value to be placed in the ``DNDarrays`` main diagonal
"""
# Todo: make this 3D/nD
if len(self.shape) != 2:
raise ValueError("Only 2D tensors supported at the moment")
if self.is_distributed():
counts, displ, _ = self.comm.counts_displs_shape(self.shape, self.split)
k = min(self.shape[0], self.shape[1])
for p in range(self.comm.size):
if displ[p] > k:
break
proc = p
if self.comm.rank <= proc:
indices = (
displ[self.comm.rank],
displ[self.comm.rank + 1] if (self.comm.rank + 1) != self.comm.size else k,
)
if self.split == 0:
self.larray[:, indices[0] : indices[1]] = self.larray[
:, indices[0] : indices[1]
].fill_diagonal_(value)
elif self.split == 1:
self.larray[indices[0] : indices[1], :] = self.larray[
indices[0] : indices[1], :
].fill_diagonal_(value)
else:
self.larray = self.larray.fill_diagonal_(value)
return self
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)
indexed_dims = len(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 value_shape[-i] != output_shape[-i] and value_shape[-i] != 1:
raise ValueError(
f"could not broadcast input array 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)
value_shape = tuple(torch.broadcast_shapes(value.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
if self.larray.numel() == 0:
return
if torch.is_tensor(key) and key.numel() == 0:
return
if key == slice(0, 0):
return
if isinstance(key, tuple) and any(
(torch.is_tensor(k) and k.numel() == 0) or k == slice(0, 0) for k in key
):
return
rhs = value.larray.type(self.dtype.torch_type()) if hasattr(value, "larray") else value
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)
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.device.torch_device
)
self.comm.Bcast(indexed_arr, root=p.root)
else:
indexed_arr = self.larray[p.key]
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_local(self, p: ProcessedKey) -> "DNDarray":
"""
Handles process-local indexing (including standard slices and local advanced indices) directly on local array partitions
without MPI communication.
"""
indexed_arr = self.larray[p.key]
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]
# 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) -> "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]
# calculate gshape
local_count = local_result.shape[0]
total_count = self.comm.allreduce(local_count, op=MPI.SUM)
gshape = (total_count,) + local_result.shape[1:]
return DNDarray(
local_result,
gshape=gshape,
dtype=self.dtype,
split=p.output_split,
device=self.device,
comm=self.comm,
balanced=False,
)
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.
"""
key = p.key
# If key was not distributed, partition it so each rank requests its share of output
if p.output_split is not None:
if isinstance(key, torch.Tensor) and key.ndim > 0:
key_split = p.output_split
if key_split < key.ndim and key.shape[key_split] == p.output_shape[p.output_split]:
k_dnd = factories.array(
key, split=key_split, comm=self.comm, device=self.device
)
key = k_dnd.larray
elif isinstance(key, tuple):
split_k = key[self.split]
if isinstance(split_k, torch.Tensor) and split_k.ndim > 0:
key_split = 0 if p.key_is_mask_like else split_k.ndim - 1
if split_k.shape[key_split] == p.output_shape[p.output_split]:
key_list = list(key)
if p.key_is_mask_like:
for idx in range(len(key_list)):
if isinstance(key_list[idx], torch.Tensor):
kd = factories.array(
key_list[idx],
split=key_split,
comm=self.comm,
device=self.device,
)
key_list[idx] = kd.larray
else:
kd = factories.array(
split_k,
split=key_split,
comm=self.comm,
device=self.device,
)
key_list[self.split] = kd.larray
key = tuple(key_list)
self, indexed_arr = self.__getitem_unordered(
key=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,
)
return indexed_arr
def __getitem_unordered(
self,
key: tuple,
output_shape: tuple,
output_split: int,
out_is_balanced: bool,
key_is_mask_like: bool,
) -> 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.empty(
sum(idx_recv_counts), dtype=split_key.dtype, device=self.device.torch_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.empty(
sum(idx_recv_counts), dtype=split_key.dtype, device=self.device.torch_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.device.torch_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 split_key.ndim > 1 and not key_is_mask_like:
original_local_shape = (
output_shape[:communication_split]
+ split_key.shape
+ output_shape[communication_split + split_key.ndim :]
)
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,
)
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]
# collect and calculate recv_counts and recv_displs
recv_counts = torch.empty_like(send_counts)
self.comm.Alltoall(send_counts, recv_counts)
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,
)
[docs]
def __getitem__(self, key: Indexer) -> DNDarray:
"""
Global getter function for DNDarrays.
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 <INDEXING>`.
Parameters
----------
key : array-like indexer
Indices to get from the ``DNDarray``.
Examples
--------
>>> a = ht.arange(10, split=0)
(1/2) >>> tensor([0, 1, 2, 3, 4], dtype=torch.int32)
(2/2) >>> tensor([5, 6, 7, 8, 9], dtype=torch.int32)
>>> a[1:6]
(1/2) >>> tensor([1, 2, 3, 4], dtype=torch.int32)
(2/2) >>> tensor([5], dtype=torch.int32)
>>> 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/2) >>> tensor([0.])
(2/2) >>> tensor([0., 0.])
"""
if key is None:
return self.expand_dims(0)
if (
key is ...
or (isinstance(key, slice) and key == slice(None))
or (isinstance(key, tuple) and key == ())
):
return self
# attempt early out for non-distributed arrays
if not self.is_distributed():
try:
res_tensor = self.larray[_unwrap_local_key(key, device=self.device.torch_device)]
return DNDarray(
res_tensor,
gshape=tuple(res_tensor.shape),
dtype=self.dtype,
split=None,
device=self.device,
comm=self.comm,
balanced=True,
)
except Exception:
pass
# 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
if op == "scalar":
return self.__getitem_scalar(processed_key)
elif op == "distr_mask":
return self.__getitem_mask(processed_key)
elif op == "distributed":
return self.__getitem_advanced_distributed(processed_key)
elif op == "descending_slice":
return self.__getitem_descending_slice_distributed(processed_key)
elif op in ("local_mask", "local"):
return self.__getitem_local(processed_key)
if torch.cuda.device_count() > 0:
def gpu(self) -> DNDarray:
"""
Returns a copy of this object in GPU memory. If this object is already in GPU memory, then no copy is
performed and the original object is returned.
"""
self.__array = self.__array.cuda(devices.gpu.torch_device)
self.__device = devices.gpu
return self
[docs]
def __index__(self) -> int:
"""
Converts a zero-dimensional integer array to a Python ``int`` object.
"""
if not issubclass(self.dtype, integer):
raise TypeError("only integer scalar arrays can be converted to a scalar index")
return self.__cast(int)
[docs]
def __int__(self) -> int:
"""
Integer scalar casting.
"""
return self.__cast(int)
[docs]
def is_balanced(self, force_check: bool = False) -> bool:
"""
Determine if ``self`` is balanced evenly (or as evenly as possible) across all nodes
distributed evenly (or as evenly as possible) across all processes.
This is equivalent to returning ``self.balanced``. If no information
is available (``self.balanced = None``), the balanced status will be
assessed via collective communication.
Parameters
----------
force_check : bool, optional
If True, the balanced status of the ``DNDarray`` will be assessed via
collective communication in any case.
"""
if not self.is_distributed():
self.__balanced = True
return self.balanced
if not force_check and self.balanced is not None:
return self.balanced
_, _, chk = self.comm.chunk(self.shape, self.split)
test_lshape = tuple([x.stop - x.start for x in chk])
balanced = 1 if test_lshape == self.lshape else 0
out = self.comm.allreduce(balanced, MPI.SUM)
balanced = True if out == self.comm.size else False
return balanced
[docs]
def is_distributed(self) -> bool:
"""
Determines whether the data of this ``DNDarray`` is distributed across multiple processes.
"""
return self.split is not None and self.comm.is_distributed()
[docs]
def item(self):
"""
Returns the only element of a 1-element :class:`DNDarray`.
Mirror of the pytorch command by the same name. If size of ``DNDarray`` is >1 element, then a ``ValueError`` is
raised (by pytorch)
Examples
--------
>>> import heat as ht
>>> x = ht.zeros((1))
>>> x.item()
0.0
"""
if self.size > 1:
raise ValueError("only one-element DNDarrays can be converted to Python scalars")
# make sure the element is on every process
self.resplit_(None)
return self.__array.item()
[docs]
def __len__(self) -> int:
"""
The length of the ``DNDarray``, i.e. the number of items in the first dimension.
"""
try:
len = self.shape[0]
return len
except IndexError:
raise TypeError("len() of unsized DNDarray")
[docs]
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.
If the ``DNDarray`` is distributed, an MPI Allgather operation will be performed before converting to np.ndarray, i.e. each MPI process will end up holding a copy of the entire array in memory. Make sure process memory is sufficient!
Examples
--------
>>> import heat as ht
T1 = ht.random.randn((10,8))
T1.numpy()
"""
dist = self.copy().resplit_(axis=None)
return dist.larray.cpu().numpy()
[docs]
def _repr_pretty_(self, p, cycle):
"""
Pretty print for IPython.
"""
if cycle:
p.text(printing.__str__(self))
else:
p.text(printing.__str__(self))
[docs]
def __repr__(self) -> str:
"""
Returns a printable representation of the passed DNDarray, targeting developers.
"""
return printing.__repr__(self)
[docs]
def ravel(self) -> DNDarray:
"""
Flattens the ``DNDarray``.
See Also
--------
:func:`~heat.core.manipulations.ravel`
Examples
--------
>>> a = ht.ones((2, 3), split=0)
>>> b = a.ravel()
>>> a[0, 0] = 4
>>> b
DNDarray([4., 1., 1., 1., 1., 1.], dtype=ht.float32, device=cpu:0, split=0)
"""
return manipulations.ravel(self)
[docs]
def redistribute_(
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``.
This is an abstraction and extension of the balance function.
Parameters
----------
lshape_map : torch.Tensor, optional
The current lshape of processes.
Units are ``[rank, lshape]``.
target_map : torch.Tensor, optional
The desired distribution across the processes.
Units are ``[rank, target lshape]``.
Note: the only important parts of the target map are the values along the split axis,
values which are not along this axis are there to mimic the shape of the ``lshape_map``.
Examples
--------
>>> st = ht.ones((50, 81, 67), split=2)
>>> target_map = torch.zeros((st.comm.size, 3), dtype=torch.int64)
>>> target_map[0, 2] = 67
>>> print(target_map)
[0/2] tensor([[ 0, 0, 67],
[0/2] [ 0, 0, 0],
[0/2] [ 0, 0, 0]], dtype=torch.int32)
[1/2] tensor([[ 0, 0, 67],
[1/2] [ 0, 0, 0],
[1/2] [ 0, 0, 0]], dtype=torch.int32)
[2/2] tensor([[ 0, 0, 67],
[2/2] [ 0, 0, 0],
[2/2] [ 0, 0, 0]], dtype=torch.int32)
>>> print(st.lshape)
[0/2] (50, 81, 23)
[1/2] (50, 81, 22)
[2/2] (50, 81, 22)
>>> st.redistribute_(target_map=target_map)
>>> print(st.lshape)
[0/2] (50, 81, 67)
[1/2] (50, 81, 0)
[2/2] (50, 81, 0)
"""
if not self.is_distributed():
return
snd_dtype = self.dtype.torch_type()
# units -> {pr, 1st index, 2nd index}
if lshape_map is None:
# NOTE: giving an lshape map which is incorrect will result in an incorrect distribution
lshape_map = self.create_lshape_map(force_check=True)
else:
if not isinstance(lshape_map, torch.Tensor):
raise TypeError(f"lshape_map must be a torch.Tensor, currently {type(lshape_map)}")
if lshape_map.shape != (self.comm.size, len(self.gshape)):
raise ValueError(
f"lshape_map must have the shape ({self.comm.size}, {len(self.gshape)}), currently {lshape_map.shape}"
)
if target_map is None: # if no target map is given then it will balance the tensor
_, _, chk = self.comm.chunk(self.shape, self.split)
target_map = lshape_map.clone()
target_map[..., self.split] = 0
for pr in range(self.comm.size):
target_map[pr, self.split] = self.comm.chunk(self.shape, self.split, rank=pr)[1][
self.split
]
self.__balanced = True
else:
sanitation.sanitize_in_tensor(target_map)
if target_map[..., self.split].sum() != self.shape[self.split]:
raise ValueError(
f"Sum along the split axis of the target map must be equal to the shape in that dimension, currently {target_map[..., self.split]}"
)
if target_map.shape != (self.comm.size, len(self.gshape)):
raise ValueError(
f"target_map must have the shape {(self.comm.size, len(self.gshape))}, currently {target_map.shape}"
)
# no info on balanced status
self.__balanced = False
lshape_cumsum = torch.cumsum(lshape_map[..., self.split], dim=0)
chunk_cumsum = torch.cat(
(
torch.tensor([0], device=self.device.torch_device),
torch.cumsum(target_map[..., self.split], dim=0),
),
dim=0,
)
# need the data start as well for process 0
for rcv_pr in range(self.comm.size - 1):
st = chunk_cumsum[rcv_pr].item()
sp = chunk_cumsum[rcv_pr + 1].item()
# start pr should be the next process with data
if lshape_map[rcv_pr, self.split] >= target_map[rcv_pr, self.split]:
# if there is more data on the process than the start process than start == stop
st_pr = rcv_pr
sp_pr = rcv_pr
else:
# if there is less data on the process than need to get the data from the next data
# with data
# need processes > rcv_pr with lshape > 0
st_pr = (
torch.nonzero(input=lshape_map[rcv_pr:, self.split] > 0, as_tuple=False)[
0
].item()
+ rcv_pr
)
hld = (
torch.nonzero(input=sp <= lshape_cumsum[rcv_pr:], as_tuple=False).flatten()
+ rcv_pr
)
sp_pr = hld[0].item() if hld.numel() > 0 else self.comm.size
# st_pr and sp_pr are the processes on which the data sits at the beginning
# need to loop from st_pr to sp_pr + 1 and send the pr
for snd_pr in range(st_pr, sp_pr + 1):
if snd_pr == self.comm.size:
break
data_required = abs(sp - st - lshape_map[rcv_pr, self.split].item())
send_amt = (
data_required
if data_required <= lshape_map[snd_pr, self.split]
else lshape_map[snd_pr, self.split]
)
if (sp - st) <= lshape_map[rcv_pr, self.split].item() or snd_pr == rcv_pr:
send_amt = 0
# send amount is the data still needed by recv if that is available on the snd
if send_amt != 0:
self.__redistribute_shuffle(
snd_pr=snd_pr, send_amt=send_amt, rcv_pr=rcv_pr, snd_dtype=snd_dtype
)
lshape_cumsum[snd_pr] -= send_amt
lshape_cumsum[rcv_pr] += send_amt
lshape_map[rcv_pr, self.split] += send_amt
lshape_map[snd_pr, self.split] -= send_amt
if lshape_map[rcv_pr, self.split] > target_map[rcv_pr, self.split]:
# if there is any data left on the process then send it to the next one
send_amt = lshape_map[rcv_pr, self.split] - target_map[rcv_pr, self.split]
self.__redistribute_shuffle(
snd_pr=rcv_pr, send_amt=send_amt.item(), rcv_pr=rcv_pr + 1, snd_dtype=snd_dtype
)
lshape_cumsum[rcv_pr] -= send_amt
lshape_cumsum[rcv_pr + 1] += send_amt
lshape_map[rcv_pr, self.split] -= send_amt
lshape_map[rcv_pr + 1, self.split] += send_amt
if any(lshape_map[..., self.split] != target_map[..., self.split]):
# sometimes need to call the redistribute once more,
# (in the case that the second to last processes needs to get data from +1 and -1)
self.redistribute_(lshape_map=lshape_map, target_map=target_map)
self.__lshape_map = target_map
self.__counts_displs = None
def __redistribute_shuffle(
self,
snd_pr: int | torch.Tensor,
send_amt: int | torch.Tensor,
rcv_pr: int | torch.Tensor,
snd_dtype: torch.dtype,
):
"""
Function to abstract the function used during redistribute for shuffling data between
processes along the split axis
Parameters
----------
snd_pr : int or torch.Tensor
Sending process
send_amt : int or torch.Tensor
Amount of data to be sent by the sending process
rcv_pr : int or torch.Tensor
Receiving process
snd_dtype : torch.dtype
Torch type of the data in question
"""
rank = self.comm.rank
send_slice = [slice(None)] * self.ndim
keep_slice = [slice(None)] * self.ndim
if rank == snd_pr:
if snd_pr < rcv_pr: # data passed to a higher rank (off the bottom)
send_slice[self.split] = slice(
self.lshape[self.split] - send_amt, self.lshape[self.split]
)
keep_slice[self.split] = slice(0, self.lshape[self.split] - send_amt)
if snd_pr > rcv_pr: # data passed to a lower rank (off the top)
send_slice[self.split] = slice(0, send_amt)
keep_slice[self.split] = slice(send_amt, self.lshape[self.split])
data = self.__array[tuple(send_slice)].clone()
self.comm.Send(data, dest=rcv_pr, tag=685)
self.__array = self.__array[tuple(keep_slice)]
if rank == rcv_pr:
shp = list(self.gshape)
shp[self.split] = send_amt
data = torch.zeros(shp, dtype=snd_dtype, device=self.device.torch_device)
self.comm.Recv(data, source=snd_pr, tag=685)
if snd_pr < rcv_pr: # data passed from a lower rank (append to top)
self.__array = torch.cat((data, self.__array), dim=self.split)
if snd_pr > rcv_pr: # data passed from a higher rank (append to bottom)
self.__array = torch.cat((self.__array, data), dim=self.split)
[docs]
def resplit_(self, axis: int = None):
"""
In-place option for resplitting a :class:`DNDarray`.
Parameters
----------
axis : int
The new split axis, ``None`` denotes gathering, an int will set the new split axis
Examples
--------
>>> a = ht.zeros(
... (
... 4,
... 5,
... ),
... split=0,
... )
>>> a.lshape
(0/2) (2, 5)
(1/2) (2, 5)
>>> ht.resplit_(a, None)
>>> a.split
None
>>> a.lshape
(0/2) (4, 5)
(1/2) (4, 5)
>>> a = ht.zeros(
... (
... 4,
... 5,
... ),
... split=0,
... )
>>> a.lshape
(0/2) (2, 5)
(1/2) (2, 5)
>>> ht.resplit_(a, 1)
>>> a.split
1
>>> a.lshape
(0/2) (4, 3)
(1/2) (4, 2)
"""
# 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 == self.split:
return self
if axis is None:
gathered = torch.empty(
self.shape, dtype=self.dtype.torch_type(), device=self.device.torch_device
)
counts, displs = self.counts_displs()
self.comm.Allgatherv(self.__array, (gathered, counts, displs), recv_axis=self.split)
self.__array = gathered
self.__split = axis
self.__lshape_map = None
self.__counts_displs = None
return self
# tensor needs be split/sliced locally
if self.split is None:
_, _, slices = self.comm.chunk(self.shape, axis)
temp = self.__array[slices]
self.__array = torch.empty((1,), device=self.device.torch_device)
# necessary to clear storage of local __array
self.__array = temp.clone().detach()
self.__split = axis
self.__lshape_map = None
self.__counts_displs = None
return self
arr_tiles = tiling.SplitTiles(self)
new_tiles = tiling.SplitTiles(self)
gshape = self.shape
new_lshape = list(gshape)
new_lshape[axis] = int(arr_tiles.tile_dimensions[axis][self.comm.rank].item())
recv_buffer = torch.empty(
tuple(new_lshape), dtype=self.dtype.torch_type(), device=self.device.torch_device
)
self._axis2axisResplit(
self.larray, self.split, arr_tiles, recv_buffer, axis, new_tiles, self.comm
)
self.__array = recv_buffer
self.__split = axis
self.__lshape_map = None
self.__counts_displs = None
return self
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:
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_local(self, p: ProcessedKey, value: "DNDarray", value_is_scalar: bool) -> None:
"""
Handles process-local item assignment (slices and local indices) directly on local partitions. If `value` is distributed, MPI communication might be necessary to align it with the target slice before assignment.
"""
if value_is_scalar:
self.__set(p.key, value)
return
value_is_distributed = isinstance(value, DNDarray) and value.is_distributed()
if not self.is_distributed() and not value_is_distributed:
self.__set(p.key, value)
return
if self.is_distributed():
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)
self.__set(p.key, value)
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.
"""
if value_is_scalar:
self.__set(p.key, value)
return
flipped_value = manipulations.flip(value, axis=p.output_split)
# determine local element count along split axis
split_key = p.key[self.split]
if isinstance(split_key, slice):
step = 1 if split_key.step is None else split_key.step
local_count = len(range(split_key.start, split_key.stop, step))
else:
local_count = split_key.numel()
# gather local slice counts across all ranks to build the target distribution map
counts = torch.empty(
(self.comm.size, 1), dtype=torch.int64, device=self.device.torch_device
)
self.comm.Allgather(
torch.tensor([local_count], dtype=torch.int64, device=self.device.torch_device),
counts,
)
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] = counts[:, 0]
flipped_value.redistribute_(target_map=target_map)
self.__set(p.key, flipped_value)
def __setitem_mask(self, p: ProcessedKey, 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 isinstance(value, (int, float, bool, complex)):
self.larray[pytorch_key] = value
return
elif hasattr(value, "larray"):
scalar_torch = value.larray.type(self.dtype.torch_type())
else:
scalar_torch = torch.as_tensor(value, device=self.device.torch_device).type(
self.dtype.torch_type()
)
self.larray[pytorch_key] = scalar_torch
else:
if isinstance(value, DNDarray) and value.is_distributed():
# value should align with local mask
value_torch = value.larray
rhs = (
value_torch
if value_torch.dtype == self.larray.dtype
else value_torch.type(self.dtype.torch_type())
)
try:
self.larray[pytorch_key] = rhs
except RuntimeError as e:
raise ValueError(
f"Shape mismatch: Cannot assign distributed array with local shape {value.lshape} "
f"on rank {self.comm.rank}: {e}"
) from e
else:
# Value is a non-distributed array -> MPI prefix sum needed
value_torch = value.larray
# 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.rank == 0:
offset = 0
_ = self.comm.exscan(local_true)
else:
offset = self.comm.exscan(local_true)
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:
# PyTorch assigns and broadcasts natively
self.larray[pytorch_key] = value_torch.type(self.dtype.torch_type())
def __setitem_advanced_distributed(
self,
p: ProcessedKey,
original_key,
value: "DNDarray",
value_is_scalar: bool,
original_split: int = None,
) -> 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()
)
value_is_distributed = isinstance(value, DNDarray) and value.is_distributed()
# 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."
)
counts, displs = self.counts_displs()
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=counts,
displs=displs,
rank=self.comm.rank,
key_is_distributed=key_is_distributed,
)
return
rank = self.comm.rank
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
)
):
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
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 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[lhs_index] = scalar_torch
return
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()
if local_indices.numel() > 0:
key_local = split_key_flat[local_indices] - displs[rank]
if value_is_scalar:
rhs = (
value.larray.type(self.dtype.torch_type())
if hasattr(value, "larray")
else value
)
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())
if self.larray.is_cuda:
key_local, rhs = _resolve_duplicate_indices(key_local, rhs, self.larray.shape)
self.larray[key_local] = rhs
return
if isinstance(original_key, tuple):
raw_split_part = original_key[original_split]
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()
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,
key_is_distributed: bool = False,
) -> 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:
if key_is_distributed:
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 = value.lshape_map
target_map[:, value.split] = global_split_key.lshape_map[:, 0]
value.redistribute_(target_map=target_map)
else:
# Key is replicated: slice locally to match value partition directly
v_counts, v_displs = value.counts_displs()
start = v_displs[rank]
end = start + v_counts[rank]
key = tuple(k[start:end] if isinstance(k, torch.Tensor) else k for k in key)
split_key = key[self.split]
else:
if key_is_distributed:
# redistribute split-axis `key` to match distribution of `value` in one pass
if key_is_single_tensor:
split_key = key
else:
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
else:
# Key is replicated: slice locally to match value partition directly
v_counts, v_displs = value.counts_displs()
start = v_displs[rank]
end = start + v_counts[rank]
if key_is_single_tensor:
key = key[start:end]
split_key = key
else:
key_list = list(key)
key_list[self.split] = key_list[self.split][start:end]
key = tuple(key_list)
split_key = key[self.split]
# 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)
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)
)
send_counts_l = send_counts.tolist()
send_displs_l = send_displs.tolist()
recv_counts_l = recv_counts.tolist()
recv_displs_l = recv_displs.tolist()
# exchange indices
if key_is_mask_like:
mask_dims = len(key)
idx_send_counts = [c * mask_dims for c in send_counts_l]
idx_send_displs = [d * mask_dims for d in send_displs_l]
idx_recv_counts = [c * mask_dims for c in recv_counts_l]
idx_recv_displs = [d * mask_dims for d in recv_displs_l]
send_idx = torch.stack([k.flatten()[sort_idx] for k in key], dim=1).reshape(-1)
recv_idx_flat = torch.empty(
sum(idx_recv_counts), dtype=split_key.dtype, device=self.device.torch_device
)
else:
idx_send_counts, idx_send_displs = send_counts_l, send_displs_l
idx_recv_counts, idx_recv_displs = recv_counts_l, recv_displs_l
send_idx = split_key_flat[sort_idx]
recv_idx_flat = torch.empty(
sum(idx_recv_counts), dtype=split_key.dtype, device=self.device.torch_device
)
self.comm.Alltoallv(
(send_idx, idx_send_counts, idx_send_displs),
(recv_idx_flat, idx_recv_counts, idx_recv_displs),
)
# exchange value
trailing_dims_shape = list(value.lshape[1:])
trailing_dims_size = math.prod(trailing_dims_shape)
val_send_counts = [c * trailing_dims_size for c in send_counts_l]
val_send_displs = [d * trailing_dims_size for d in send_displs_l]
val_recv_counts = [c * trailing_dims_size for c in recv_counts_l]
val_recv_displs = [d * trailing_dims_size for d in recv_displs_l]
send_vals = value.larray[sort_idx].contiguous().reshape(-1)
recv_vals_flat = torch.empty(
sum(val_recv_counts), dtype=value.larray.dtype, device=self.device.torch_device
)
self.comm.Alltoallv(
(send_vals, val_send_counts, val_send_displs),
(recv_vals_flat, val_recv_counts, val_recv_displs),
)
if key_is_mask_like:
recv_indices = recv_idx_flat.reshape(sum(recv_counts_l), len(key))
recv_indices[:, 0] -= displs[rank]
key = tuple(recv_indices[:, i] for i in range(len(key)))
else:
# store incoming indices in int 1-D tensor and correct for rank offset
recv_indices = recv_idx_flat - displs[rank]
# 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)
recv_vals = recv_vals_flat.reshape(-1, *trailing_dims_shape)
recv_buf = DNDarray(
recv_vals.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)
return self
[docs]
def __setitem__(
self,
key: Indexer,
value: float | "DNDarray" | torch.Tensor,
):
"""
Global item setter for DNDarrays.
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 <INDEXING>`.
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.]])
"""
if not self.is_distributed() and not (
isinstance(value, DNDarray) and value.is_distributed()
):
try:
torch_key = _unwrap_local_key(key, device=self.device.torch_device)
if isinstance(value, DNDarray):
rhs = value.larray.to(self.larray.dtype)
elif isinstance(value, torch.Tensor):
rhs = value.to(self.larray.dtype)
else:
rhs = value # Python scalar / float / int
if self.larray.is_cuda and torch.is_tensor(rhs):
torch_key, rhs = _resolve_duplicate_indices(torch_key, rhs, self.larray.shape)
self.larray[torch_key] = rhs
return
except Exception:
pass
# bypass factories.array() for primitive types and single-element tensors to avoid unnecessary overhead
value_is_primitive = isinstance(value, (int, float, complex, bool)) or (
isinstance(value, torch.Tensor) and value.numel() == 1 and value.ndim == 0
)
if not value_is_primitive and not isinstance(value, DNDarray):
value = factories.array(value)
original_key = key
original_split = self.split
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:
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, 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, original_split=original_split
)
elif op == "descending_slice":
self.__setitem_descending_slice_distributed(processed_key, value, value_is_scalar)
elif op in ("local_mask", "local"):
self.__setitem_local(processed_key, value, value_is_scalar)
[docs]
def __str__(self) -> str:
"""
Computes a string representation of the passed ``DNDarray``.
"""
return printing.__str__(self)
[docs]
def to_device(self, device: Device, /, *, stream: int | Any | None = None) -> DNDarray:
"""
Copy the array from the device on which it currently resides to the specified ``device``.
Parameters
----------
device : Device
A ``Device`` object.
stream : Int or Any, optional
Stream object to use during copy.
"""
if stream is not None:
raise ValueError("The stream argument to to_device() is not supported")
if device.device_type == "cpu":
return self.cpu()
elif device.device_type == "gpu":
return self.gpu()
raise ValueError(f"Unsupported device {device!r}")
[docs]
def tolist(self, keepsplit: bool = False) -> list[int | float]:
"""
Return a copy of the local array data as a (nested) Python list. For scalars, a standard Python number is returned.
Parameters
----------
keepsplit: bool
Whether the list should be returned locally or globally.
Examples
--------
>>> a = ht.array([[0, 1], [2, 3]])
>>> a.tolist()
[[0, 1], [2, 3]]
>>> a = ht.array([[0, 1], [2, 3]], split=0)
>>> a.tolist()
[[0, 1], [2, 3]]
>>> a = ht.array([[0, 1], [2, 3]], split=1)
>>> a.tolist(keepsplit=True)
(1/2) [[0], [2]]
(2/2) [[1], [3]]
"""
if not keepsplit:
return self.resplit(axis=None).__array.tolist()
return self.__array.tolist()
[docs]
@classmethod
def __torch_function__(cls, func, types, args=(), kwargs=None):
"""
Supports PyTorch's dispatch mechanism.
"""
import heat
if kwargs is None:
kwargs = {}
try:
ht_func = getattr(heat, func.__name__)
except AttributeError:
return NotImplemented
return ht_func(*args, **kwargs)
[docs]
def __torch_proxy__(self) -> torch.Tensor:
"""
Return a 1-element `torch.Tensor` strided as the global `self` shape.
Used internally for sanitation purposes.
"""
return torch.empty(self.gshape, device="meta")
# Heat imports at the end to break cyclic dependencies
from . import complex_math
from . import devices
from . import factories
from . import indexing
from . import linalg
from . import manipulations
from . import printing
from . import rounding
from . import sanitation
from . import statistics
from . import stride_tricks
from . import tiling
from . import types
from .devices import Device
from .stride_tricks import sanitize_axis
from .types import datatype, integer, canonical_heat_type
from .types import bool as ht_bool, uint8 as ht_uint8
