Features/914 adv indexing

heat/core/dndarray.py

@@ -9,7 +9,7 @@
 from inspect import stack
 from mpi4py import MPI
 from pathlib import Path
-from typing import List, Union, Tuple, TypeVar, Optional
+from typing import List, Union, Tuple, TypeVar, Optional, Iterable
 
 warnings.simplefilter("always", ResourceWarning)
 
@@ -653,6 +653,84 @@ def fill_diagonal(self, value: float) -> DNDarray:
 
         return self
 
+    def __process_key(arr: DNDarray, key: Union[int, Tuple[int, ...], List[int, ...]]) -> Tuple:
+        """
+        Private method for processing keys for indexing. Returns wether advanced indexing is used as well as a processed key and self.
+        A processed key:
+        - doesn't cotain any ellipses or newaxis
+        - all Iterables are converted to ``DNDarrays``
+        - has the same dimensionality as the ``DNDarray`` it indexes
+
+        Parameters
+        ----------
+        key : int, slice, Tuple[int,...], List[int,...]
+            Indices for the tensor.
+        """
+        advanced_indexing = False
+        if isinstance(key, DNDarray):
+            # DNDARRAY CURRENTLY DOES NOT IMPLEMENT the Iterable interface, need to define __iter__()
+            advanced_indexing = True
+            # TODO: check for key.ndim = 0 and treat that as int
+            # TODO: get outshape + outsplit; depends on wether key is bool or int and key.ndim
+        elif isinstance(key, Iterable) and not isinstance(key, tuple):
+            advanced_indexing = True
+            key = factories.array(key)
+            # DOES NOT WORK FOR SEQUENCE OF TENSORS OR DNDARRAYS, works for sequence of ndarrays though
+            # TODO: get outshape + outsplit; depends on wether key is bool or int and key.ndim
+        elif isinstance(key, tuple):
+            key = list(key)
+            for i, k in enumerate(key):
+                if isinstance(k, Iterable) or isinstance(key, DNDarray):
+                    advanced_indexing = True
+                    key[i] = factories.array(key[i])
+                    # DOES NOT WORK FOR SEQUENCE OF TENSORS OR DNDARRAYS, works for seq of ndarrays though
+                    # TODO: check for key.ndim = 0 and treat that as int
+                    # TODO: get outshape + outsplit; depends on wether key is bool or int and key.ndim
+            add_dims = sum(k is None for k in key)  # (np.newaxis is None)===true
+            ellipsis = sum(isinstance(k, type(...)) for k in key)
+            if ellipsis > 1:
+                raise ValueError("key can only contain 1 ellipsis")
+            elif ellipsis == 1:
+                expand_key = [slice(None)] * (arr.ndim + add_dims)
+                ellipsis_index = key.index(...)
+                expand_key[:ellipsis_index] = key[:ellipsis_index]
+                expand_key[ellipsis_index - len(key) :] = key[ellipsis_index + 1 :]
+                key = expand_key
+            if add_dims:
+                for i, k in reversed(enumerate(key)):
+                    if k is None:
+                        key[i] = slice(None)
+                        arr = arr.expand_dims(i - add_dims + 1)  # is the -1 correct?
+                        add_dims -= 1
+            # expand key to match the number of dimensions of the DNDarray
+            key = tuple(key + [slice(None)] * (arr.ndim - len(key)))
+        else:  # key is integer or slice
+            key = tuple([key] + [slice(None)] * (arr.ndim - 1))
+        return advanced_indexing, arr, key
+
+    def __get_local_slice(self, key: slice):
+        split = self.split
+        if split is None:
+            return key
+        key = stride_tricks.sanitize_slice(key, self.shape[split])
+        start, stop, step = key.start, key.stop, key.step
+        if step < 0:  # NOT supported by torch, should be filtered by torch_proxy
+            key = self.__get_local_slice(slice(stop + 1, start + 1, abs(step)))
+            if key is None:
+                return None
+            start, stop, step = key.start, key.stop, key.step
+            return slice(key.stop - 1, key.start - 1, -1 * key.step)
+
+        _, offsets = self.counts_displs()
+        offset = offsets[self.comm.rank]
+        range_proxy = range(self.lshape[split])
+        local_inds = range_proxy[start - offset : stop - offset]  # only works if stop - offset > 0
+        local_inds = local_inds[max(offset - start, 0) % step :: step]
+        if len(local_inds) and stop > offset:
+            # otherwise if (stop-offset) > -self.lshape[split] this can index into the local chunk despite ending before it
+            return slice(local_inds.start, local_inds.stop, local_inds.step)
+        return None
+
     def __getitem__(self, key: Union[int, Tuple[int, ...], List[int, ...]]) -> DNDarray:
         """
         Global getter function for DNDarrays.
@@ -684,65 +762,54 @@ def __getitem__(self, key: Union[int, Tuple[int, ...], List[int, ...]]) -> DNDar
         (2/2) >>> tensor([0., 0.])
         """
         # key can be: int, tuple, list, slice, DNDarray, torch tensor, numpy array, or sequence thereof
-        self_proxy = self.__torch_proxy__()
-        self_proxy.names = [
-            "split" if (self.split is not None and i == self.split) else "_{}".format(i)
-            for i in range(self_proxy.ndim)
+        # Trivial cases
+        if key is None:
+            return self.expand_dims(0)
+        if key == ... or key == slice(None):  # latter doesnt work with torch for 0-dim tensors
+            return self
+        # Preprocess: Process Ellipsis + 'None' indexing; make Iterables to DNDarrays
+        advanced_indexing, self, key = self.__process_key(key)
+
+        # To use torch_proxy with advanced indexing, add empty dimensions instead of
+        # advanced index. Later, replace the empty dimensions with the shape of the advanced index
+        proxy = self
+        names = [
+            "split" if (proxy.split is not None and i == proxy.split) else "_{}".format(i)
+            for i in range(proxy.ndim)
         ]
+        proxy_key = list(key)
+        if advanced_indexing:
+            proxy_key = list(key)
+            for i, k in reversed(enumerate(key)):
+                if isinstance(k, DNDarray):  # all iterables have been made DNDarrays
+                    # TODO: Bool indexing (sometimes) is collapsed into one dimension
+                    # TODO: What to do if advanced index is in split dimension??
+                    names[i] = "replace" + str(k.shape)  # put shape into name
+                    proxy_key[i] = slice(None)
+                    for _ in range(k.ndim - 1):
+                        proxy = proxy.expand_dims(i)
+                        names.insert(i + 1, "_{}".format(len(names)))
+                        proxy_key.insert(i + 1, slice(None))
+        proxy_key = tuple(proxy_key)
+
+        self_proxy = proxy.__torch_proxy__()
+        self_proxy.names = names
+        indexed_proxy = self_proxy[proxy_key]
+
+        output_shape = list(indexed_proxy.shape)
+        if advanced_indexing:
+            for i, n in enumerate(indexed_proxy.names):
+                if "replace" in n:
+                    shape = eval(n.split("replace")[1])  # extract shape from name
+                    # TODO Bool indexing (sometimes) is collapsed into one dimension
+                    output_shape[i : i + len(shape)] = shape
+        output_shape = tuple(output_shape)
 
-        try:
-            indexed_proxy = self_proxy[key]
-        except IndexError as e:
-            # key might be a DNDarray or contain DNDarrays, torch returns IndexError
-            try:
-                # key might be a DNDarray
-                key_proxy = key.__torch_proxy__()
-                key_proxy.names = [
-                    "split" if (key.split is not None and i == key.split) else "_{}".format(i)
-                    for i in range(key_proxy.ndim)
-                ]
-                indexed_proxy = self_proxy[key_proxy]
-            except AttributeError:
-                # key might be sequence of DNDarrays
-                key = list(key.copy())
-                for i in len(key):
-                    if isinstance(key[i], DNDarray):
-                        if key[i].is_distributed:
-                            raise NotImplementedError(
-                                "Advanced indexing with distributed DNDarrays not supported yet"
-                            )
-                        key[i] = key[i].larray
-                try:
-                    indexed_proxy = self_proxy[tuple(key)]
-                except IndexError:
-                    raise e
-        # TODO: catch torch exceptions, return reasonable error message
-
-        output_shape = tuple(indexed_proxy.shape)
         try:
             output_split = indexed_proxy.names.index("split")
         except ValueError:
             output_split = None
 
-        try:
-            key_ndims = getattr(key, "ndim", len(key))
-        except TypeError:
-            # key is a scalar or a slice
-            key = (key,)
-            key_ndims = 1
-
-        # expand key to match the number of dimensions of the DNDarray
-        if key_ndims < self.ndim:
-            expand_key = [slice(None)] * self.ndim
-            # account for ellipsis
-            if key.count(...):
-                ellipsis_index = key.index(...)
-                expand_key[:ellipsis_index] = key[:ellipsis_index]
-                expand_key[ellipsis_index + 2 :] = key[ellipsis_index + 1 :]
-            else:
-                expand_key[:key_ndims] = key
-            key = tuple(expand_key)
-
         # data are not distributed or split dimension is not affected by indexing
         if not self.is_distributed or key[self.split] == slice(None):
             return DNDarray(
@@ -758,79 +825,14 @@ def __getitem__(self, key: Union[int, Tuple[int, ...], List[int, ...]]) -> DNDar
         # data are distributed and split dimension is affected by indexing
         _, offsets = self.counts_displs()
         split = self.split
-
         # slice along the split axis
         if isinstance(key[split], slice):
-            if key[split].start is None:
-                slice_start = 0
-            else:
-                slice_start = (
-                    key[split].start
-                    if key[split].start > 0
-                    else key[split].start + self.gshape[split]
-                )
-            if key[split].stop is None:
-                slice_stop = self.gshape[split]
-            else:
-                slice_stop = (
-                    key[split].stop if key[split].stop > 0 else key[split].stop + self.gshape[split]
-                )
-            slice_step = key[split].step
-
-            # identify active ranks
-            offsets = torch.tensor(offsets, dtype=torch.int64, device=self.larray.device)
-            first_active = torch.where(offsets - slice_start <= 0)[0][-1].item()
-            last_active = torch.where(offsets - slice_stop <= 0)[0][-1].item()
-            active_ranks = range(first_active, last_active + 1)
-
-            if self.comm.rank in active_ranks:
-                if slice_step is None:
-                    slice_step = 1
-                # calculate local slice
-                if (
-                    slice_start >= offsets[self.comm.rank]
-                    and slice_start < self.lshape[split] + offsets[self.comm.rank]
-                ):
-                    local_slice_start = slice_start - offsets[self.comm.rank]
-                else:
-                    if slice_step != 1:
-                        local_slice_start = torch.arange(
-                            offsets[self.comm.rank],
-                            offsets[self.comm.rank] + slice_step,
-                            dtype=torch.int64,
-                            device=self.larray.device,
-                        )
-                        local_slice_start = (
-                            torch.where(local_slice_start % slice_step == 0)[0].item()
-                            - offsets[self.comm.rank]
-                        )
-                    else:
-                        local_slice_start = 0
-                if (
-                    slice_stop >= offsets[self.comm.rank]
-                    and slice_stop < self.lshape[split] + offsets[self.comm.rank]
-                ):
-                    local_slice_stop = slice_stop - offsets[self.comm.rank]
-                else:
-                    if slice_step != 1:
-                        local_slice_stop = torch.arange(
-                            offsets[self.comm.rank] + 1 - slice_step,
-                            offsets[self.comm.rank] + 1,
-                            dtype=torch.int64,
-                            device=self.larray.device,
-                        )
-                        local_slice_stop = (
-                            torch.where(local_slice_stop % slice_step == 0)[0].item()
-                            - offsets[self.comm.rank]
-                        )
-                    else:
-                        local_slice_stop = self.lshape[split]
-                # slice local tensor
-                local_slice = slice(local_slice_start, local_slice_stop, slice_step)
-                key = key[:split] + (local_slice,) + key[split + 1 :]
-                local_tensor = self.larray[key]
-            else:
-                # local tensor is empty
+            local_slice = self.__get_local_slice(key[split])
+            if local_slice is not None:
+                key = list(key)
+                key[split] = local_slice
+                local_tensor = self.larray[tuple(key)]
+            else:  # local tensor is empty
                 local_shape = list(output_shape)
                 local_shape[output_split] = 0
                 local_tensor = torch.zeros(
@@ -1569,6 +1571,40 @@ def __setitem__(
         (2/2) >>> tensor([[0., 1., 0., 0., 0.],
                           [0., 1., 0., 0., 0.]])
         """
+
+        def __set(arr: DNDarray, value: DNDarray):
+            """
+            Setter for not advanced indexing, i.e. when arr[key] is an in-place view of arr.
+            """
+            if not isinstance(value, DNDarray):
+                value = factories.array(value, device=arr.device, comm=arr.comm)
+            while value.ndim < arr.ndim:  # broadcasting
+                value = value.expand_dims(0)
+            sanitation.sanitize_out(arr, value.shape, value.split, value.device, value.comm)
+            value = sanitation.sanitize_distribution(value, target=arr)
+            arr.larray[None] = value.larray
+            return
+
+        if key is None or key == ... or key == slice(None):
+            return __set(self, value)
+
+        advanced_indexing, self, key = self.__process_key(key)
+        if advanced_indexing:
+            raise Exception("Advanced indexing is not supported yet")
+
+        split = self.split
+        if not self.is_distributed or key[split] == slice(None):
+            return __set(self[key], value)
+
+        if isinstance(key[split], slice):
+            return __set(self[key], value)
+
+        if np.isscalar(key[split]):
+            key = list(key)
+            idx = int(key[split])
+            key[split] = slice(idx, idx + 1)
+            return __set(self[tuple(key)], value)
+
         key = getattr(key, "copy()", key)
         try:
             if value.split != self.split: