Bug/825 setitem slice dndarrays

CHANGELOG.md

@@ -50,6 +50,7 @@ Example on 2 processes:
 - [#820](https://github.com/helmholtz-analytics/heat/pull/820) `randn` values are pushed away from 0 by the minimum value the given dtype before being transformed into the Gaussian shape
 - [#821](https://github.com/helmholtz-analytics/heat/pull/821) Fixed `__getitem__` handling of distributed `DNDarray` key element
 - [#826](https://github.com/helmholtz-analytics/heat/pull/826) Fixed `__setitem__` handling of distributed `DNDarray` values which have a different shape in the split dimension
+- [#831](https://github.com/helmholtz-analytics/heat/pull/831) `__getitem__` handling of `array-like` 1-element key
 
 ## Feature additions
 ### Exponential

heat/core/dndarray.py

@@ -576,18 +576,18 @@ def cpu(self) -> DNDarray:
         self.__device = devices.cpu
         return self
 
-    def create_lshape_map(self, recreate: bool = True) -> torch.Tensor:
+    def create_lshape_map(self, force_check: bool = True) -> torch.Tensor:
         """
         Generate a 'map' of the lshapes of the data on all processes.
         Units are ``(process rank, lshape)``
 
         Parameters
         ----------
-        recreate : bool, optional
+        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 recreate and self.__lshape_map is not None:
+        if not force_check and self.__lshape_map is not None:
             return self.__lshape_map
 
         lshape_map = torch.zeros(
@@ -1367,6 +1367,18 @@ def __setitem__(
                           [0., 1., 0., 0., 0.]])
         """
         key = getattr(key, "copy()", key)
+        try:
+            if value.split != self.split:
+                val_split = int(value.split)
+                sp = self.split
+                warnings.warn(
+                    f"\nvalue.split {val_split} not equal to this DNDarray's split:"
+                    f" {sp}. this may cause errors or unwanted behavior",
+                    category=RuntimeWarning,
+                )
+        except (AttributeError, TypeError):
+            pass
+
         if isinstance(key, DNDarray) and key.ndim == self.ndim:
             # this splits the key into torch.Tensors in each dimension for advanced indexing
             lkey = [slice(None, None, None)] * self.ndim
@@ -1392,6 +1404,13 @@ def __setitem__(
             kend = key[ell_ind + 1 :]
             slices = [slice(None)] * (self.ndim - (len(kst) + len(kend)))
             key = kst + slices + kend
+
+        for c, k in enumerate(key):
+            try:
+                key[c] = k.item()
+            except (AttributeError, ValueError):
+                pass
+
         key = tuple(key)
 
         if not self.is_distributed():
@@ -1411,6 +1430,8 @@ def __setitem__(
         chunk_start = chunk_slice[self.split].start
         chunk_end = chunk_slice[self.split].stop
 
+        self_proxy = torch.ones((1,)).as_strided(self.gshape, [0] * self.ndim)
+
         if not isinstance(key, tuple):
             return self.__setter(key, value)  # returns None
 
@@ -1448,7 +1469,6 @@ def __setitem__(
                 target_reshape_map = torch.zeros(
                     (self.comm.size, self.ndim), dtype=torch.int, device=self.device.torch_device
                 )
-                self_proxy = torch.ones((1,)).as_strided(self.gshape, [0] * self.ndim)
                 for r in range(self.comm.size):
                     if r not in actives:
                         loc_key = key.copy()

heat/core/tests/test_dndarray.py

@@ -1089,7 +1089,7 @@ def test_setitem_getitem(self):
 
         # slice in 1st dim across 1 node (2nd) w/ singular second dim
         c = ht.zeros((13, 5), split=0)
-        c[8:12, 1] = 1
+        c[8:12, ht.array(1)] = 1
         b = c[8:12, np.int64(1)]
         self.assertTrue((b == 1).all())
         self.assertEqual(b.gshape, (4,))