Small refactor of the use of nonzero in _resolve_indexing_state

Author: brownbaerchen

heat/core/dndarray.py

@@ -358,12 +358,7 @@ def _resolve_indexing_state(
             and first_dtype in (ht_bool, ht_uint8, torch.bool, torch.uint8, np.bool_, np.uint8)
         ):
             if isinstance(first, DNDarray):
-                nz = first.nonzero()
-                if isinstance(nz, tuple):
-                    nz = nz[0]
-                if getattr(nz, "ndim", 1) > 1 and nz.shape[-1] == 1:
-                    nz = nz.squeeze(-1)
-                idx0 = nz
+                idx0 = first.nonzero(as_tuple=False).flatten()
             elif isinstance(first, torch.Tensor):
                 idx0 = torch.nonzero(first, as_tuple=False).flatten()
             else:  # np.ndarray

heat/core/indexing.py

@@ -146,7 +146,7 @@ def nonzero(x: DNDarray, as_tuple: bool = True) -> tuple[DNDarray, ...] | DNDarr
     )
 
 
-DNDarray.nonzero = lambda self: nonzero(self, as_tuple=True)
+DNDarray.nonzero = nonzero
 DNDarray.nonzero.__doc__ = nonzero.__doc__