Repository: helmholtz-analytics/heat

Issue:
`ht.nonzero()` should return tuple of 1-D arrays instead of an n-D array
**Description**
`np.nonzero(a)` returns a tuple of 1-D ndarrays with length `a.ndim`. This can be then used directly and without modifications for advanced indexing across dimensions.
```
>>> a
array([[0, 1, 2],
       [3, 4, 5],
       [6, 7, 8]])
>>> mask
array([[False,  True,  True],
       [False,  True,  True],
       [False,  True,  True]])
>>> a[mask].shape
(6,)
>>> a[np.nonzero(mask)].shape
(6,)
>>> a[mask] == a[np.nonzero(mask)]
array([ True,  True,  True,  True,  True,  True])
```
`ht.nonzero` returns the indices in `torch.nonzero` format, in the example above it returns a DNDarray with shape (6,2). 

This is not what the users expect, and complicates our advanced indexing unnecessarily.

Will be addressed together with #914 


**Expected behavior**
`ht.nonzero` should return a tuple of 1-D dndarrays, for API consistency and to simplify `heat`'s advanced indexing implementation

**Version Info**
`main`




Discussion:

Resolving patch:
diff --git a/heat/core/dndarray.py b/heat/core/dndarray.py
--- a/heat/core/dndarray.py
+++ b/heat/core/dndarray.py
@@ -872,7 +872,7 @@ def __getitem__(self, key: Union[int, Tuple[int, ...], List[int, ...]]) -> DNDar
             output_split = None
 
         # data are not distributed or split dimension is not affected by indexing
-        if not self.is_distributed or key[self.split] == slice(None):
+        if not self.is_distributed() or key[self.split] == slice(None):
             return DNDarray(
                 self.larray[key],
                 gshape=output_shape,
@@ -1654,7 +1654,7 @@ def __set(arr: DNDarray, value: DNDarray):
             raise Exception("Advanced indexing is not supported yet")
 
         split = self.split
-        if not self.is_distributed or key[split] == slice(None):
+        if not self.is_distributed() or key[split] == slice(None):
             return __set(self[key], value)
 
         if isinstance(key[split], slice):
diff --git a/heat/core/indexing.py b/heat/core/indexing.py
--- a/heat/core/indexing.py
+++ b/heat/core/indexing.py
@@ -13,12 +13,12 @@
 __all__ = ["nonzero", "where"]
 
 
-def nonzero(x: DNDarray) -> DNDarray:
+def nonzero(x: DNDarray) -> Tuple[DNDarray, ...]:
     """
-    Return a :class:`~heat.core.dndarray.DNDarray` containing the indices of the elements that are non-zero.. (using ``torch.nonzero``)
-    If ``x`` is split then the result is split in the 0th dimension. However, this :class:`~heat.core.dndarray.DNDarray`
+    Return a Tuple of :class:`~heat.core.dndarray.DNDarray`s, one for each dimension of ``x``,
+    containing the indices of the non-zero elements in that dimension. If ``x`` is split then
+    the result is split in the 0th dimension. However, this :class:`~heat.core.dndarray.DNDarray`
     can be UNBALANCED as it contains the indices of the non-zero elements on each node.
-    Returns an array with one entry for each dimension of ``x``, containing the indices of the non-zero elements in that dimension.
     The values in ``x`` are always tested and returned in row-major, C-style order.
     The corresponding non-zero values can be obtained with: ``x[nonzero(x)]``.
 
@@ -32,10 +32,8 @@ def nonzero(x: DNDarray) -> DNDarray:
     >>> import heat as ht
     >>> x = ht.array([[3, 0, 0], [0, 4, 1], [0, 6, 0]], split=0)
     >>> ht.nonzero(x)
-    DNDarray([[0, 0],
-              [1, 1],
-              [1, 2],
-              [2, 1]], dtype=ht.int64, device=cpu:0, split=0)
+    (DNDarray([0, 1, 1, 2], dtype=ht.int64, device=cpu:0, split=None),
+        DNDarray([0, 1, 2, 1], dtype=ht.int64, device=cpu:0, split=None))
     >>> y = ht.array([[1, 2, 3], [4, 5, 6], [7, 8, 9]], split=0)
     >>> y > 3
     DNDarray([[False, False, False],
@@ -48,6 +46,8 @@ def nonzero(x: DNDarray) -> DNDarray:
               [2, 0],
               [2, 1],
               [2, 2]], dtype=ht.int64, device=cpu:0, split=0)
+    (DNDarray([1, 1, 1, 2, 2, 2], dtype=ht.int64, device=cpu:0, split=None),
+        DNDarray([0, 1, 2, 0, 1, 2], dtype=ht.int64, device=cpu:0, split=None))
     >>> y[ht.nonzero(y > 3)]
     DNDarray([4, 5, 6, 7, 8, 9], dtype=ht.int64, device=cpu:0, split=0)
     """
@@ -56,39 +56,42 @@ def nonzero(x: DNDarray) -> DNDarray:
     except AttributeError:
         raise TypeError("Input must be a DNDarray, is {}".format(type(x)))
 
+    lcl_nonzero = torch.nonzero(input=local_x, as_tuple=False).transpose(0, 1)
+
     if x.split is None:
-        # if there is no split then just return the values from torch
-        lcl_nonzero = torch.nonzero(input=local_x, as_tuple=False)
+        # if there is no split then just return the transpose of values from torch
+
         gout = list(lcl_nonzero.size())
         is_split = None
     else:
         # a is split
-        lcl_nonzero = torch.nonzero(input=local_x, as_tuple=False)
         # adjust local indices along split dimension
         _, displs = x.counts_displs()
-        lcl_nonzero[..., x.split] += displs[x.comm.rank]
+        lcl_nonzero[x.split] += displs[x.comm.rank]
         del displs
+
         # get global size of split dimension
         gout = list(lcl_nonzero.size())
-        gout[0] = x.comm.allreduce(gout[0], MPI.SUM)
+        gout[1] = x.comm.allreduce(gout[1], MPI.SUM)
         is_split = 0
 
-    if x.ndim == 1:
-        lcl_nonzero = lcl_nonzero.squeeze(dim=1)
-    for g in range(len(gout) - 1, -1, -1):
-        if gout[g] == 1:
-            del gout[g]
-
-    return DNDarray(
-        lcl_nonzero,
-        gshape=tuple(gout),
-        dtype=types.canonical_heat_type(lcl_nonzero.dtype),
-        split=is_split,
-        device=x.device,
-        comm=x.comm,
-        balanced=False,
+    non_zero_indices = list(
+        [
+            DNDarray(
+                dim_indices,
+                gshape=tuple(gout),
+                dtype=types.canonical_heat_type(lcl_nonzero.dtype),
+                split=is_split,
+                device=x.device,
+                comm=x.comm,
+                balanced=False,
+            )
+            for dim_indices in lcl_nonzero
+        ]
     )
 
+    return tuple(non_zero_indices)
+
 
 DNDarray.nonzero = lambda self: nonzero(self)
 DNDarray.nonzero.__doc__ = nonzero.__doc__


Test patch:
diff --git a/heat/core/tests/test_indexing.py b/heat/core/tests/test_indexing.py
--- a/heat/core/tests/test_indexing.py
+++ b/heat/core/tests/test_indexing.py
@@ -9,18 +9,18 @@ def test_nonzero(self):
         a = ht.array([[1, 2, 3], [4, 5, 2], [7, 8, 9]], split=None)
         cond = a > 3
         nz = ht.nonzero(cond)
-        self.assertEqual(nz.gshape, (5, 2))
-        self.assertEqual(nz.dtype, ht.int64)
-        self.assertEqual(nz.split, None)
+        self.assertEqual(len(nz), 2)
+        self.assertEqual(len(nz[0]), 5)
+        self.assertEqual(nz[0].dtype, ht.int64)
 
         # split
         a = ht.array([[1, 2, 3], [4, 5, 6], [7, 8, 9]], split=1)
         cond = a > 3
         nz = cond.nonzero()
-        self.assertEqual(nz.gshape, (6, 2))
-        self.assertEqual(nz.dtype, ht.int64)
-        self.assertEqual(nz.split, 0)
-        a[nz] = 10.0
+        self.assertEqual(len(nz), 2)
+        self.assertEqual(len(nz[0]), 6)
+        self.assertEqual(nz[0].dtype, ht.int64)
+        a[nz] = 10
         self.assertEqual(ht.all(a[nz] == 10), 1)
 
     def test_where(self):