Indexing with 0-dimensional arrays

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

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)