impl iterative print corners

Author: vtavana

dpctl/tensor/_print.py

@@ -15,6 +15,7 @@
 #  limitations under the License.
 
 import contextlib
+import itertools
 import operator
 
 import numpy as np
@@ -236,20 +237,28 @@ def _nd_corners(arr_in, edge_items):
         sycl_queue=arr_in.sycl_queue,
     )
 
+    split_dim = sum(
+        arr_in.shape[dim] > 2 * edge_items for dim in range(arr_ndim)
+    )
+    blocks = [
+        list(ele) for ele in list(itertools.product(*[[0, 1]] * split_dim))
+    ]
+    for dim in range(arr_ndim):
+        if arr_in.shape[dim] <= 2 * edge_items:
+            for blk in blocks:
+                blk.insert(dim, 0)
+
     hev_list = []
-    for corner in range(arr_ndim**2):
+    for blk in blocks:
         slices = ()
-        tmp = bin(corner).replace("0b", "").zfill(arr_ndim)
-
         for dim in reversed(range(arr_ndim)):
-            if arr_in.shape[dim] < 2 * edge_items:
-                slices = (np.s_[:],) + slices
-            else:
-                ind = (-1) ** int(tmp[dim]) * edge_items
-                if ind < 0:
+            if arr_in.shape[dim] > 2 * edge_items:
+                if blk[dim] == 1:
                     slices = (np.s_[-edge_items::],) + slices
                 else:
                     slices = (np.s_[:edge_items:],) + slices
+            else:
+                slices = (np.s_[:],) + slices
         hev, _ = ti._copy_usm_ndarray_into_usm_ndarray(
             src=arr_in[slices],
             dst=arr_out[slices],