diff --git a/heat/core/dndarray.py b/heat/core/dndarray.py index 9ec0ea89e1..ee4774781b 100644 --- a/heat/core/dndarray.py +++ b/heat/core/dndarray.py @@ -690,7 +690,7 @@ def __getitem__(self, key: Union[int, Tuple[int, ...], List[int, ...]]) -> DNDar # NOTE: this gathers the entire key on every process!! # TODO: remove this resplit!! key = manipulations.resplit(key) - if key.larray.dtype in [torch.bool, torch.uint8]: + if key.larray.dtype in [torch.bool, torch.uint8] and key.ndim > 0: key = indexing.nonzero(key) if key.ndim > 1: @@ -703,12 +703,12 @@ def __getitem__(self, key: Union[int, Tuple[int, ...], List[int, ...]]) -> DNDar advanced_ind = True elif not isinstance(key, tuple): """this loop handles all other cases. DNDarrays which make it to here refer to - advanced indexing slices, as do the torch tensors. Both DNDaarrys and torch.Tensors + advanced indexing slices, as do the torch tensors. Both DNDarrays and torch.Tensors are cast into lists here by PyTorch. lists mean advanced indexing will be used""" h = [slice(None, None, None)] * max(self.ndim, 1) if isinstance(key, DNDarray): key = manipulations.resplit(key) - if key.larray.dtype in [torch.bool, torch.uint8]: + if key.larray.dtype in [torch.bool, torch.uint8] and key.ndim > 0: h[0] = torch.nonzero(key.larray).flatten() # .tolist() else: h[0] = key.larray.tolist() @@ -965,7 +965,14 @@ def __key_is_singular(key: any, axis: int, self_proxy: torch.Tensor) -> bool: def __key_adds_dimension(key: any, axis: int, self_proxy: torch.Tensor) -> bool: # determine if the key adds a new dimension zeros = (0,) * (self_proxy.ndim - 1) - return self_proxy[(*zeros[:axis], key[axis], *zeros[axis:])].ndim == 2 + if key[axis] is None: + indexed_dims = self_proxy.shape[:-1] + else: + indexed_dims = self_proxy.shape + if 0 in indexed_dims or len(indexed_dims) > self_proxy.ndim: + return False + indexed_proxy = self_proxy[(*zeros[:axis], key[axis], *zeros[axis:])] + return 0 not in indexed_proxy.shape and indexed_proxy.ndim == 2 def item(self): """ @@ -1414,7 +1421,7 @@ def __setitem__( into the torch tensors for each dimension. This signals that advanced indexing is to be used.""" key = manipulations.resplit(key) - if key.larray.dtype in [torch.bool, torch.uint8]: + if key.larray.dtype in [torch.bool, torch.uint8] and key.ndim > 0: key = indexing.nonzero(key) if key.ndim > 1: @@ -1426,9 +1433,9 @@ def __setitem__( key = [key] elif not isinstance(key, tuple): """this loop handles all other cases. DNDarrays which make it to here refer to - advanced indexing slices, as do the torch tensors. Both DNDaarrys and torch.Tensors + advanced indexing slices, as do the torch tensors. Both DNDarrays and torch.Tensors are cast into lists here by PyTorch. lists mean advanced indexing will be used""" - h = [slice(None, None, None)] * self.ndim + h = [slice(None, None, None)] * max(self.ndim, 1) if isinstance(key, DNDarray): key = manipulations.resplit(key) if key.larray.dtype in [torch.bool, torch.uint8]: @@ -1456,7 +1463,7 @@ def __setitem__( pass # remove bools from a torch tensor in favor of indexes try: - if key[i].dtype in [torch.bool, torch.uint8]: + if key[i].dtype in [torch.bool, torch.uint8] and key[i].ndim > 0: key[i] = torch.nonzero(key[i]).flatten() except (AttributeError, TypeError): pass diff --git a/heat/core/tests/test_dndarray.py b/heat/core/tests/test_dndarray.py index e42c5a9a14..51b7f1705a 100644 --- a/heat/core/tests/test_dndarray.py +++ b/heat/core/tests/test_dndarray.py @@ -1327,6 +1327,10 @@ def test_setitem_getitem(self): if a.comm.rank == 1: self.assertEqual(a[10, ...].lshape, (5, 3)) + a = ht.array(False) + a[...] = False + self.assertEqual(a, ht.array(False)) + a = ht.zeros((13, 5, 8), split=2) # # set value on one node a[10, 0, 0] = 1 @@ -1407,6 +1411,15 @@ def test_setitem_getitem(self): self.assertTrue(np.all(arr.numpy() == np_arr)) self.assertTrue(ht.all(arr[ht_key] == 10.0)) + arr = ht.array(False) + arr_indexed_bool = arr[ht.array(False)] + self.assertTrue(arr_indexed_bool.shape == (0,)) + self.assertTrue(arr_indexed_bool.size == 0) + arr[ht.array(False)] = False + self.assertEqual(arr, ht.array(False)) + arr = ht.array([False]) + self.assertTrue(arr[ht.array(False)].shape == (0, 1)) + split = 0 arr = ht.random.random((20, 20)).resplit(split) np_arr = arr.numpy() @@ -1493,6 +1506,9 @@ def test_setitem_getitem(self): self.assertTrue(arr.shape == check.shape) self.assertTrue(arr.lshape[new_dim] == 1) + empty = ht.array([]) + self.assertTrue(empty[None].shape == (1, 0)) + def test_size_gnumel(self): a = ht.zeros((10, 10, 10), split=None) self.assertEqual(a.size, 10 * 10 * 10)