From b4eafff64ed4dcf0bf64818d64279a9c26268101 Mon Sep 17 00:00:00 2001
From: Vahid Tavanashad <vahid.tavanashad@intel.com>
Date: Mon, 24 Apr 2023 16:53:44 -0500
Subject: [PATCH 1/2] initial commit

---
 dpctl/tensor/_print.py | 56 ++++++++++++++++++++++++++++--------------
 1 file changed, 38 insertions(+), 18 deletions(-)

diff --git a/dpctl/tensor/_print.py b/dpctl/tensor/_print.py
index ce6923597e..8401230c91 100644
--- a/dpctl/tensor/_print.py
+++ b/dpctl/tensor/_print.py
@@ -19,7 +19,9 @@
 
 import numpy as np
 
+import dpctl
 import dpctl.tensor as dpt
+import dpctl.tensor._tensor_impl as ti
 
 __doc__ = "Print functions for :class:`dpctl.tensor.usm_ndarray`."
 
@@ -220,25 +222,43 @@ def print_options(*args, **kwargs):
         dpt.set_print_options(**options)
 
 
-def _nd_corners(x, edge_items, slices=()):
-    axes_reduced = len(slices)
-    if axes_reduced == x.ndim:
-        return x[slices]
-
-    if x.shape[axes_reduced] > 2 * edge_items:
-        return dpt.concat(
-            (
-                _nd_corners(
-                    x, edge_items, slices + (slice(None, edge_items, None),)
-                ),
-                _nd_corners(
-                    x, edge_items, slices + (slice(-edge_items, None, None),)
-                ),
-            ),
-            axis=axes_reduced,
+def _nd_corners(arr_in, edge_items):
+    arr_ndim = arr_in.ndim
+    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)
+    )
+
+    arr_out = dpt.empty(
+        res_shape,
+        dtype=arr_in.dtype,
+        usm_type=arr_in.usm_type,
+        sycl_queue=arr_in.sycl_queue,
+    )
+
+    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
+        hev, _ = ti._copy_usm_ndarray_into_usm_ndarray(
+            src=arr_in[slices],
+            dst=arr_out[slices],
+            sycl_queue=arr_in.sycl_queue,
         )
-    else:
-        return _nd_corners(x, edge_items, slices + (slice(None, None, None),))
+        hev_list.append(hev)
+
+    dpctl.SyclEvent.wait_for(hev_list)
+    return arr_out
 
 
 def usm_ndarray_str(

From e153d77b2ed2fbf835440595376ba56fb03e2eea Mon Sep 17 00:00:00 2001
From: Vahid Tavanashad <vahid.tavanashad@intel.com>
Date: Tue, 25 Apr 2023 12:52:59 -0500
Subject: [PATCH 2/2] impl iterative print corners

---
 dpctl/tensor/_print.py | 45 +++++++++++++++++++++---------------------
 1 file changed, 23 insertions(+), 22 deletions(-)

diff --git a/dpctl/tensor/_print.py b/dpctl/tensor/_print.py
index 8401230c91..30bb353df7 100644
--- a/dpctl/tensor/_print.py
+++ b/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,13 @@ 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
+    if max(_shape) <= max_shape:
+        return dpt.asnumpy(arr_in)
     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 +240,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 +367,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)