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
@@ -223,10 +224,11 @@ def print_options(*args, **kwargs):
 
 
 def _nd_corners(arr_in, edge_items):
-    arr_ndim = arr_in.ndim
+    _shape = arr_in.shape
+    max_shape = 2 * edge_items + 1
     res_shape = tuple(
-        2 * edge_items if arr_in.shape[i] > 2 * edge_items else arr_in.shape[i]
-        for i in range(arr_ndim)
+        max_shape if _shape[i] > max_shape else _shape[i]
+        for i in range(arr_in.ndim)
     )
 
     arr_out = dpt.empty(
@@ -236,29 +238,27 @@ def _nd_corners(arr_in, edge_items):
         sycl_queue=arr_in.sycl_queue,
     )
 
+    blocks = []
+    for i in range(len(_shape)):
+        if _shape[i] > max_shape:
+            blocks.append(
+                (
+                    np.s_[:edge_items],
+                    np.s_[-edge_items:],
+                )
+            )
+        else:
+            blocks.append((np.s_[:],))
+
     hev_list = []
-    for corner in range(arr_ndim**2):
-        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:
-                    slices = (np.s_[-edge_items::],) + slices
-                else:
-                    slices = (np.s_[:edge_items:],) + slices
+    for slc in itertools.product(*blocks):
         hev, _ = ti._copy_usm_ndarray_into_usm_ndarray(
-            src=arr_in[slices],
-            dst=arr_out[slices],
-            sycl_queue=arr_in.sycl_queue,
+            src=arr_in[slc], dst=arr_out[slc], sycl_queue=arr_in.sycl_queue
         )
         hev_list.append(hev)
 
     dpctl.SyclEvent.wait_for(hev_list)
-    return arr_out
+    return dpt.asnumpy(arr_out)
 
 
 def usm_ndarray_str(
@@ -365,8 +365,7 @@ def usm_ndarray_str(
     edge_items = options["edgeitems"]
 
     if x.size > threshold:
-        # need edge_items + 1 elements for np.array2string to abbreviate
-        data = dpt.asnumpy(_nd_corners(x, edge_items + 1))
+        data = _nd_corners(x, edge_items)
         options["threshold"] = 0
     else:
         data = dpt.asnumpy(x)