From 6641d1eb607d081aaa5ef8d2e44621a2b887c7eb Mon Sep 17 00:00:00 2001 From: Ben Bourgart Date: Tue, 22 Feb 2022 14:31:44 +0100 Subject: [PATCH 1/3] Preprocess key, workaround torch_proxy for advanced indexing, simplify slice-indexing. UNTESTED --- heat/core/dndarray.py | 215 ++++++++++++++++++------------------------ 1 file changed, 94 insertions(+), 121 deletions(-) diff --git a/heat/core/dndarray.py b/heat/core/dndarray.py index 2786583ebe..3c069dc93c 100644 --- a/heat/core/dndarray.py +++ b/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) @@ -684,65 +684,93 @@ 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__() + # 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 = 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)] * (self.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) + self = self.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)] * (self.ndim - len(key))) + else: # key is integer or slice + key = tuple([key] + [slice(None)] * (self.ndim - 1)) + + # 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_key = key + proxy = self + if advanced_indexing: + proxy_key = [] + replace = {} + 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 + replace[i] = k.shape + proxy_key.extend([slice(None)] * k.ndim) + for _ in range(k.ndim - 1): + proxy = proxy.expand_dims(i) + else: + proxy_key.append(k) + proxy_key = tuple(reversed(proxy_key)) + + self_proxy = proxy.__torch_proxy__() self_proxy.names = [ - "split" if (self.split is not None and i == self.split) else "_{}".format(i) + "split" if (proxy.split is not None and i == proxy.split) else "_{}".format(i) for i in range(self_proxy.ndim) ] + indexed_proxy = self_proxy[proxy_key] - 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 = list(indexed_proxy.shape) + if advanced_indexing: + for i, shape in replace.values(): + # TODO Bool indexing (sometimes) is collapsed into one dimension + output_shape[i : i + len(shape)] = shape + output_shape = tuple(output_shape) - 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 +786,24 @@ 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 + key = list(key) + key[split] = stride_tricks.sanitize_slice(key[split], self.shape[split]) + start, stop, step = key[split].start, key[split].stop, key[split].step + if step < 0: # NOT supported by torch; TODO throw Exception + key[split] = slice(stop + 1, start + 1, abs(step)) + return self[tuple(key)].flip(axis=self.split) + + offset = offsets[self.comm.rank] + range_proxy = range(self.lshape[split]) + local_inds = range_proxy[start - offset : stop - offset] + local_inds = local_inds[(offset - start) % step :: step] + if len(local_inds): + local_slice = slice(local_inds.start, local_inds.stop, local_inds.step) + 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( From cd78ecbbe67fb2b6ece25ea0e7fad4f03abe2cb3 Mon Sep 17 00:00:00 2001 From: Ben Bourgart Date: Tue, 22 Feb 2022 15:46:41 +0100 Subject: [PATCH 2/3] put advanced index shape in the dimensions name to get the correct position in the index_proxy --- heat/core/dndarray.py | 54 +++++++++++++++++++++---------------------- 1 file changed, 27 insertions(+), 27 deletions(-) diff --git a/heat/core/dndarray.py b/heat/core/dndarray.py index 3c069dc93c..3d32ead45f 100644 --- a/heat/core/dndarray.py +++ b/heat/core/dndarray.py @@ -691,26 +691,23 @@ def __getitem__(self, key: Union[int, Tuple[int, ...], List[int, ...]]) -> DNDar return self # Preprocess: Process Ellipsis + 'None' indexing; make Iterables to DNDarrays advanced_indexing = False - if isinstance( - key, DNDarray - ): # DNDARRAY CURRENTLY DOES NOT IMPLEMENT the Iterable interface, need to define __iter__() + 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 + 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 + 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 @@ -736,34 +733,37 @@ def __getitem__(self, key: Union[int, Tuple[int, ...], List[int, ...]]) -> DNDar # 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_key = key 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 = [] - replace = {} + 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 - replace[i] = k.shape - proxy_key.extend([slice(None)] * k.ndim) + # 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) - else: - proxy_key.append(k) - proxy_key = tuple(reversed(proxy_key)) + 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 = [ - "split" if (proxy.split is not None and i == proxy.split) else "_{}".format(i) - for i in range(self_proxy.ndim) - ] + self_proxy.names = names indexed_proxy = self_proxy[proxy_key] output_shape = list(indexed_proxy.shape) if advanced_indexing: - for i, shape in replace.values(): - # TODO Bool indexing (sometimes) is collapsed into one dimension - output_shape[i : i + len(shape)] = shape + 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: @@ -791,14 +791,14 @@ def __getitem__(self, key: Union[int, Tuple[int, ...], List[int, ...]]) -> DNDar key = list(key) key[split] = stride_tricks.sanitize_slice(key[split], self.shape[split]) start, stop, step = key[split].start, key[split].stop, key[split].step - if step < 0: # NOT supported by torch; TODO throw Exception + if step < 0: # NOT supported by torch, should be filtered by torch_proxy key[split] = slice(stop + 1, start + 1, abs(step)) return self[tuple(key)].flip(axis=self.split) offset = offsets[self.comm.rank] range_proxy = range(self.lshape[split]) local_inds = range_proxy[start - offset : stop - offset] - local_inds = local_inds[(offset - start) % step :: step] + local_inds = local_inds[max(offset - start, 0) % step :: step] if len(local_inds): local_slice = slice(local_inds.start, local_inds.stop, local_inds.step) key[split] = local_slice From 7d97ea2cd765fb394819cf4dd98c23a2bd238d40 Mon Sep 17 00:00:00 2001 From: Ben Bourgart Date: Tue, 22 Feb 2022 23:04:11 +0100 Subject: [PATCH 3/3] first changes to setitem --- heat/core/dndarray.py | 159 +++++++++++++++++++++++++++++------------- 1 file changed, 111 insertions(+), 48 deletions(-) diff --git a/heat/core/dndarray.py b/heat/core/dndarray.py index 3d32ead45f..4743cc3b4b 100644 --- a/heat/core/dndarray.py +++ b/heat/core/dndarray.py @@ -653,43 +653,19 @@ def fill_diagonal(self, value: float) -> DNDarray: return self - def __getitem__(self, key: Union[int, Tuple[int, ...], List[int, ...]]) -> DNDarray: + def __process_key(arr: DNDarray, key: Union[int, Tuple[int, ...], List[int, ...]]) -> Tuple: """ - Global getter function for DNDarrays. - Returns a new DNDarray composed of the elements of the original tensor selected by the indices - given. This does *NOT* redistribute or rebalance the resulting tensor. If the selection of values is - unbalanced then the resultant tensor is also unbalanced! - To redistributed the ``DNDarray`` use :func:`balance()` (issue #187) + 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 to get from the tensor. - - 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.]) + Indices for the tensor. """ - # key can be: int, tuple, list, slice, DNDarray, torch tensor, numpy array, or sequence thereof - # 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 = False if isinstance(key, DNDarray): # DNDARRAY CURRENTLY DOES NOT IMPLEMENT the Iterable interface, need to define __iter__() @@ -715,7 +691,7 @@ def __getitem__(self, key: Union[int, Tuple[int, ...], List[int, ...]]) -> DNDar if ellipsis > 1: raise ValueError("key can only contain 1 ellipsis") elif ellipsis == 1: - expand_key = [slice(None)] * (self.ndim + add_dims) + 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 :] @@ -724,12 +700,75 @@ def __getitem__(self, key: Union[int, Tuple[int, ...], List[int, ...]]) -> DNDar for i, k in reversed(enumerate(key)): if k is None: key[i] = slice(None) - self = self.expand_dims(i - add_dims + 1) # is the -1 correct? + 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)] * (self.ndim - len(key))) + key = tuple(key + [slice(None)] * (arr.ndim - len(key))) else: # key is integer or slice - key = tuple([key] + [slice(None)] * (self.ndim - 1)) + 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. + Returns a new DNDarray composed of the elements of the original tensor selected by the indices + given. This does *NOT* redistribute or rebalance the resulting tensor. If the selection of values is + unbalanced then the resultant tensor is also unbalanced! + To redistributed the ``DNDarray`` use :func:`balance()` (issue #187) + + Parameters + ---------- + key : int, slice, Tuple[int,...], List[int,...] + Indices to get from the tensor. + + 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.]) + """ + # key can be: int, tuple, list, slice, DNDarray, torch tensor, numpy array, or sequence thereof + # 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 @@ -788,19 +827,9 @@ def __getitem__(self, key: Union[int, Tuple[int, ...], List[int, ...]]) -> DNDar split = self.split # slice along the split axis if isinstance(key[split], slice): - key = list(key) - key[split] = stride_tricks.sanitize_slice(key[split], self.shape[split]) - start, stop, step = key[split].start, key[split].stop, key[split].step - if step < 0: # NOT supported by torch, should be filtered by torch_proxy - key[split] = slice(stop + 1, start + 1, abs(step)) - return self[tuple(key)].flip(axis=self.split) - - offset = offsets[self.comm.rank] - range_proxy = range(self.lshape[split]) - local_inds = range_proxy[start - offset : stop - offset] - local_inds = local_inds[max(offset - start, 0) % step :: step] - if len(local_inds): - local_slice = slice(local_inds.start, local_inds.stop, local_inds.step) + 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 @@ -1542,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: