Move ti.broadcast_to() to dpctl_ext/tensor and reuse it in dpctl_ext/tensor

Author: vlad-perevezentsev

dpctl_ext/tensor/__init__.py

@@ -63,6 +63,7 @@
     take_along_axis,
 )
 from dpctl_ext.tensor._manipulation_functions import (
+    broadcast_to,
     repeat,
     roll,
 )
@@ -76,6 +77,7 @@
     "asarray",
     "asnumpy",
     "astype",
+    "broadcast_to",
     "can_cast",
     "copy",
     "clip",

dpctl_ext/tensor/_clip.py

@@ -205,9 +205,9 @@ def _clip_none(x, val, out, order, _binary_fn):
                     order=order,
                 )
         if x_shape != res_shape:
-            x = dpt.broadcast_to(x, res_shape)
+            x = dpt_ext.broadcast_to(x, res_shape)
         if val_ary.shape != res_shape:
-            val_ary = dpt.broadcast_to(val_ary, res_shape)
+            val_ary = dpt_ext.broadcast_to(val_ary, res_shape)
         _manager = SequentialOrderManager[exec_q]
         dep_evs = _manager.submitted_events
         ht_binary_ev, binary_ev = _binary_fn(
@@ -251,8 +251,8 @@ def _clip_none(x, val, out, order, _binary_fn):
                 )
 
         if x_shape != res_shape:
-            x = dpt.broadcast_to(x, res_shape)
-        buf = dpt.broadcast_to(buf, res_shape)
+            x = dpt_ext.broadcast_to(x, res_shape)
+        buf = dpt_ext.broadcast_to(buf, res_shape)
         ht_binary_ev, binary_ev = _binary_fn(
             src1=x,
             src2=buf,
@@ -580,11 +580,11 @@ def clip(x, /, min=None, max=None, out=None, order="K"):
                         order=order,
                     )
             if x_shape != res_shape:
-                x = dpt.broadcast_to(x, res_shape)
+                x = dpt_ext.broadcast_to(x, res_shape)
             if a_min.shape != res_shape:
-                a_min = dpt.broadcast_to(a_min, res_shape)
+                a_min = dpt_ext.broadcast_to(a_min, res_shape)
             if a_max.shape != res_shape:
-                a_max = dpt.broadcast_to(a_max, res_shape)
+                a_max = dpt_ext.broadcast_to(a_max, res_shape)
             _manager = SequentialOrderManager[exec_q]
             dep_ev = _manager.submitted_events
             ht_binary_ev, binary_ev = ti._clip(
@@ -639,10 +639,10 @@ def clip(x, /, min=None, max=None, out=None, order="K"):
                         order=order,
                     )
 
-            x = dpt.broadcast_to(x, res_shape)
+            x = dpt_ext.broadcast_to(x, res_shape)
             if a_min.shape != res_shape:
-                a_min = dpt.broadcast_to(a_min, res_shape)
-            buf2 = dpt.broadcast_to(buf2, res_shape)
+                a_min = dpt_ext.broadcast_to(a_min, res_shape)
+            buf2 = dpt_ext.broadcast_to(buf2, res_shape)
             ht_binary_ev, binary_ev = ti._clip(
                 src=x,
                 min=a_min,
@@ -695,10 +695,10 @@ def clip(x, /, min=None, max=None, out=None, order="K"):
                         order=order,
                     )
 
-            x = dpt.broadcast_to(x, res_shape)
-            buf1 = dpt.broadcast_to(buf1, res_shape)
+            x = dpt_ext.broadcast_to(x, res_shape)
+            buf1 = dpt_ext.broadcast_to(buf1, res_shape)
             if a_max.shape != res_shape:
-                a_max = dpt.broadcast_to(a_max, res_shape)
+                a_max = dpt_ext.broadcast_to(a_max, res_shape)
             ht_binary_ev, binary_ev = ti._clip(
                 src=x,
                 min=buf1,
@@ -766,9 +766,9 @@ def clip(x, /, min=None, max=None, out=None, order="K"):
                     order=order,
                 )
 
-        x = dpt.broadcast_to(x, res_shape)
-        buf1 = dpt.broadcast_to(buf1, res_shape)
-        buf2 = dpt.broadcast_to(buf2, res_shape)
+        x = dpt_ext.broadcast_to(x, res_shape)
+        buf1 = dpt_ext.broadcast_to(buf1, res_shape)
+        buf2 = dpt_ext.broadcast_to(buf2, res_shape)
         ht_, clip_ev = ti._clip(
             src=x,
             min=buf1,

dpctl_ext/tensor/_copy_utils.py

@@ -368,7 +368,7 @@ def _put_multi_index(ary, inds, p, vals, mode=0):
         rhs = vals
     else:
         rhs = dpt_ext.astype(vals, ary.dtype)
-    rhs = dpt.broadcast_to(rhs, expected_vals_shape)
+    rhs = dpt_ext.broadcast_to(rhs, expected_vals_shape)
     _manager = dpctl.utils.SequentialOrderManager[exec_q]
     dep_ev = _manager.submitted_events
     hev, put_ev = ti._put(

dpctl_ext/tensor/_ctors.py

@@ -1200,7 +1200,7 @@ def full(
             usm_type=usm_type,
             sycl_queue=sycl_queue,
         )
-        return dpt_ext.copy(dpt.broadcast_to(X, shape), order=order)
+        return dpt_ext.copy(dpt_ext.broadcast_to(X, shape), order=order)
     else:
         _validate_fill_value(fill_value)
 
@@ -1307,7 +1307,7 @@ def full_like(
                 usm_type=usm_type,
                 sycl_queue=sycl_queue,
             )
-            X = dpt.broadcast_to(X, sh)
+            X = dpt_ext.broadcast_to(X, sh)
             res = _empty_like_orderK(x, dtype, usm_type, sycl_queue)
             _manager = dpctl.utils.SequentialOrderManager[sycl_queue]
             # order copy after tasks populating X

dpctl_ext/tensor/_indexing_functions.py

@@ -341,7 +341,7 @@ def put_vec_duplicates(vec, ind, vals):
         rhs = vals
     else:
         rhs = dpt_ext.astype(vals, x.dtype)
-    rhs = dpt.broadcast_to(rhs, val_shape)
+    rhs = dpt_ext.broadcast_to(rhs, val_shape)
 
     _manager = dpctl.utils.SequentialOrderManager[exec_q]
     deps_ev = _manager.submitted_events

dpctl_ext/tensor/_manipulation_functions.py

@@ -86,6 +86,60 @@ def _broadcast_shape_impl(shapes):
     return tuple(common_shape)
 
 
+def _broadcast_strides(X_shape, X_strides, res_ndim):
+    """
+    Broadcasts strides to match the given dimensions;
+    returns tuple type strides.
+    """
+    out_strides = [0] * res_ndim
+    X_shape_len = len(X_shape)
+    str_dim = -X_shape_len
+    for i in range(X_shape_len):
+        shape_value = X_shape[i]
+        if not shape_value == 1:
+            out_strides[str_dim] = X_strides[i]
+        str_dim += 1
+
+    return tuple(out_strides)
+
+
+def broadcast_to(X, /, shape):
+    """broadcast_to(x, shape)
+
+    Broadcast an array to a new `shape`; returns the broadcasted
+    :class:`dpctl.tensor.usm_ndarray` as a view.
+
+    Args:
+        x (usm_ndarray): input array
+        shape (Tuple[int,...]): array shape. The `shape` must be
+            compatible with `x` according to broadcasting rules.
+
+    Returns:
+        usm_ndarray:
+            An array with the specified `shape`.
+            The output array is a view of the input array, and
+            hence has the same data type, USM allocation type and
+            device attributes.
+    """
+    if not isinstance(X, dpt.usm_ndarray):
+        raise TypeError(f"Expected usm_ndarray type, got {type(X)}.")
+
+    # Use numpy.broadcast_to to check the validity of the input
+    # parameter 'shape'. Raise ValueError if 'X' is not compatible
+    # with 'shape' according to NumPy's broadcasting rules.
+    new_array = np.broadcast_to(
+        np.broadcast_to(np.empty(tuple(), dtype="u1"), X.shape), shape
+    )
+    new_sts = _broadcast_strides(X.shape, X.strides, new_array.ndim)
+    return dpt.usm_ndarray(
+        shape=new_array.shape,
+        dtype=X.dtype,
+        buffer=X,
+        strides=new_sts,
+        offset=X._element_offset,
+    )
+
+
 def repeat(x, repeats, /, *, axis=None):
     """repeat(x, repeats, axis=None)
 

dpctl_ext/tensor/_search_functions.py

@@ -389,11 +389,11 @@ def where(condition, x1, x2, /, *, order="K", out=None):
             )
 
     if condition_shape != res_shape:
-        condition = dpt.broadcast_to(condition, res_shape)
+        condition = dpt_ext.broadcast_to(condition, res_shape)
     if x1_shape != res_shape:
-        x1 = dpt.broadcast_to(x1, res_shape)
+        x1 = dpt_ext.broadcast_to(x1, res_shape)
     if x2_shape != res_shape:
-        x2 = dpt.broadcast_to(x2, res_shape)
+        x2 = dpt_ext.broadcast_to(x2, res_shape)
 
     dep_evs = _manager.submitted_events
     hev, where_ev = ti._where(