Move ti.broadcast_arrays() to dpctl_ext.tensor and reuse it

Author: vlad-perevezentsev

dpctl_ext/tensor/__init__.py

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

dpctl_ext/tensor/_copy_utils.py

@@ -306,7 +306,7 @@ def _prepare_indices_arrays(inds, q, usm_type):
     )
 
     # broadcast
-    inds = dpt.broadcast_arrays(*inds)
+    inds = dpt_ext.broadcast_arrays(*inds)
 
     return inds
 

dpctl_ext/tensor/_ctors.py

@@ -1513,7 +1513,7 @@ def meshgrid(*arrays, indexing="xy"):
         res.append(dpt_ext.reshape(array, sh, copy=True))
         sh = sh[-1:] + sh[:-1]
 
-    output = dpt.broadcast_arrays(*res)
+    output = dpt_ext.broadcast_arrays(*res)
 
     return output
 

dpctl_ext/tensor/_manipulation_functions.py

@@ -47,6 +47,15 @@
 )
 
 
+def _broadcast_shapes(*args):
+    """
+    Broadcast the input shapes into a single shape;
+    returns tuple broadcasted shape.
+    """
+    array_shapes = [array.shape for array in args]
+    return _broadcast_shape_impl(array_shapes)
+
+
 def _broadcast_shape_impl(shapes):
     if len(set(shapes)) == 1:
         return shapes[0]
@@ -103,6 +112,37 @@ def _broadcast_strides(X_shape, X_strides, res_ndim):
     return tuple(out_strides)
 
 
+def broadcast_arrays(*args):
+    """broadcast_arrays(*arrays)
+
+    Broadcasts one or more :class:`dpctl.tensor.usm_ndarrays` against
+    one another.
+
+    Args:
+        arrays (usm_ndarray): an arbitrary number of arrays to be
+            broadcasted.
+
+    Returns:
+        List[usm_ndarray]:
+            A list of broadcasted arrays. Each array
+            must have the same shape. Each array must have the same `dtype`,
+            `device` and `usm_type` attributes as its corresponding input
+            array.
+    """
+    if len(args) == 0:
+        raise ValueError("`broadcast_arrays` requires at least one argument")
+    for X in args:
+        if not isinstance(X, dpt.usm_ndarray):
+            raise TypeError(f"Expected usm_ndarray type, got {type(X)}.")
+
+    shape = _broadcast_shapes(*args)
+
+    if all(X.shape == shape for X in args):
+        return args
+
+    return [broadcast_to(X, shape) for X in args]
+
+
 def broadcast_to(X, /, shape):
     """broadcast_to(x, shape)
 

dpnp/dpnp_iface_arraycreation.py

@@ -3131,7 +3131,7 @@ def meshgrid(*xi, copy=True, sparse=False, indexing="xy"):
         output[1] = dpt_ext.reshape(output[1], (-1, 1) + s0[2:])
 
     if not sparse:
-        output = dpt.broadcast_arrays(*output)
+        output = dpt_ext.broadcast_arrays(*output)
 
     if copy:
         output = [dpt_ext.copy(x) for x in output]

dpnp/dpnp_iface_indexing.py

@@ -260,7 +260,7 @@ def choose(a, choices, out=None, mode="wrap"):
                 choices,
             )
         )
-    arrs_broadcast = dpt.broadcast_arrays(inds, *choices)
+    arrs_broadcast = dpt_ext.broadcast_arrays(inds, *choices)
     inds = arrs_broadcast[0]
     choices = tuple(arrs_broadcast[1:])