Merge pull request #798 from IntelPython/broadcast_arrays

Implementation of broadcast_arrays function

Author: oleksandr-pavlyk

dpctl/tensor/__init__.py

@@ -26,6 +26,7 @@
 from dpctl.tensor._device import Device
 from dpctl.tensor._dlpack import from_dlpack
 from dpctl.tensor._manipulation_functions import (
+    broadcast_arrays,
     broadcast_to,
     expand_dims,
     permute_dims,
@@ -42,6 +43,7 @@
     "copy",
     "empty",
     "reshape",
+    "broadcast_arrays",
     "broadcast_to",
     "expand_dims",
     "permute_dims",

dpctl/tensor/_manipulation_functions.py

@@ -15,6 +15,8 @@
 #  limitations under the License.
 
 
+from itertools import chain, repeat
+
 import numpy as np
 from numpy.core.numeric import normalize_axis_tuple
 
@@ -38,6 +40,47 @@ def _broadcast_strides(X_shape, X_strides, res_ndim):
     return tuple(out_strides)
 
 
+def _broadcast_shapes(*args):
+    """
+    Broadcast the input shapes into a single shape;
+    returns tuple broadcasted shape.
+    """
+    shapes = [array.shape for array in args]
+    if len(set(shapes)) == 1:
+        return shapes[0]
+    mutable_shapes = False
+    nds = [len(s) for s in shapes]
+    biggest = max(nds)
+    for i in range(len(args)):
+        diff = biggest - nds[i]
+        if diff > 0:
+            ty = type(shapes[i])
+            shapes[i] = ty(chain(repeat(1, diff), shapes[i]))
+    common_shape = []
+    for axis in range(biggest):
+        lengths = [s[axis] for s in shapes]
+        unique = set(lengths + [1])
+        if len(unique) > 2:
+            raise ValueError(
+                "Shape mismatch: two or more arrays have "
+                f"incompatible dimensions on axis ({axis},)"
+            )
+        elif len(unique) == 2:
+            unique.remove(1)
+            new_length = unique.pop()
+            common_shape.append(new_length)
+            for i in range(len(args)):
+                if shapes[i][axis] == 1:
+                    if not mutable_shapes:
+                        shapes = [list(s) for s in shapes]
+                        mutable_shapes = True
+                    shapes[i][axis] = new_length
+        else:
+            common_shape.append(1)
+
+    return tuple(common_shape)
+
+
 def permute_dims(X, axes):
     """
     permute_dims(X: usm_ndarray, axes: tuple or list) -> usm_ndarray
@@ -146,3 +189,21 @@ def broadcast_to(X, shape):
         strides=new_sts,
         offset=X.__sycl_usm_array_interface__.get("offset", 0),
     )
+
+
+def broadcast_arrays(*args):
+    """
+    broadcast_arrays(*args: usm_ndarrays) -> list of usm_ndarrays
+
+    Broadcasts one or more usm_ndarrays against one another.
+    """
+    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]

dpctl/tests/test_usm_ndarray_manipulation.py

@@ -326,3 +326,158 @@ def test_broadcast_to_raises(data):
     Xnp = np.zeros(orig_shape)
     X = dpt.asarray(Xnp, sycl_queue=q)
     pytest.raises(ValueError, dpt.broadcast_to, X, target_shape)
+
+
+def assert_broadcast_correct(input_shapes):
+    try:
+        q = dpctl.SyclQueue()
+    except dpctl.SyclQueueCreationError:
+        pytest.skip("Queue could not be created")
+    np_arrays = [np.zeros(s) for s in input_shapes]
+    out_np_arrays = np.broadcast_arrays(*np_arrays)
+    usm_arrays = [dpt.asarray(Xnp, sycl_queue=q) for Xnp in np_arrays]
+    out_usm_arrays = dpt.broadcast_arrays(*usm_arrays)
+    for Xnp, X in zip(out_np_arrays, out_usm_arrays):
+        assert_array_equal(
+            Xnp, dpt.asnumpy(X), err_msg=f"Failed for {input_shapes})"
+        )
+
+
+def assert_broadcast_arrays_raise(input_shapes):
+    try:
+        q = dpctl.SyclQueue()
+    except dpctl.SyclQueueCreationError:
+        pytest.skip("Queue could not be created")
+    usm_arrays = [dpt.asarray(np.zeros(s), sycl_queue=q) for s in input_shapes]
+    pytest.raises(ValueError, dpt.broadcast_arrays, *usm_arrays)
+
+
+def test_broadcast_arrays_same():
+    try:
+        q = dpctl.SyclQueue()
+    except dpctl.SyclQueueCreationError:
+        pytest.skip("Queue could not be created")
+    Xnp = np.arange(10)
+    Ynp = np.arange(10)
+    res_Xnp, res_Ynp = np.broadcast_arrays(Xnp, Ynp)
+    X = dpt.asarray(Xnp, sycl_queue=q)
+    Y = dpt.asarray(Ynp, sycl_queue=q)
+    res_X, res_Y = dpt.broadcast_arrays(X, Y)
+    assert_array_equal(res_Xnp, dpt.asnumpy(res_X))
+    assert_array_equal(res_Ynp, dpt.asnumpy(res_Y))
+
+
+def test_broadcast_arrays_one_off():
+    try:
+        q = dpctl.SyclQueue()
+    except dpctl.SyclQueueCreationError:
+        pytest.skip("Queue could not be created")
+    Xnp = np.array([[1, 2, 3]])
+    Ynp = np.array([[1], [2], [3]])
+    res_Xnp, res_Ynp = np.broadcast_arrays(Xnp, Ynp)
+    X = dpt.asarray(Xnp, sycl_queue=q)
+    Y = dpt.asarray(Ynp, sycl_queue=q)
+    res_X, res_Y = dpt.broadcast_arrays(X, Y)
+    assert_array_equal(res_Xnp, dpt.asnumpy(res_X))
+    assert_array_equal(res_Ynp, dpt.asnumpy(res_Y))
+
+
+@pytest.mark.parametrize(
+    "shapes",
+    [
+        (),
+        (1,),
+        (3,),
+        (0, 1),
+        (0, 3),
+        (1, 0),
+        (3, 0),
+        (1, 3),
+        (3, 1),
+        (3, 3),
+    ],
+)
+def test_broadcast_arrays_same_shapes(shapes):
+    for shape in shapes:
+        single_input_shapes = [shape]
+        assert_broadcast_correct(single_input_shapes)
+        double_input_shapes = [shape, shape]
+        assert_broadcast_correct(double_input_shapes)
+        triple_input_shapes = [shape, shape, shape]
+        assert_broadcast_correct(triple_input_shapes)
+
+
+@pytest.mark.parametrize(
+    "shapes",
+    [
+        [[(1,), (3,)]],
+        [[(1, 3), (3, 3)]],
+        [[(3, 1), (3, 3)]],
+        [[(1, 3), (3, 1)]],
+        [[(1, 1), (3, 3)]],
+        [[(1, 1), (1, 3)]],
+        [[(1, 1), (3, 1)]],
+        [[(1, 0), (0, 0)]],
+        [[(0, 1), (0, 0)]],
+        [[(1, 0), (0, 1)]],
+        [[(1, 1), (0, 0)]],
+        [[(1, 1), (1, 0)]],
+        [[(1, 1), (0, 1)]],
+    ],
+)
+def test_broadcast_arrays_same_len_shapes(shapes):
+    # Check that two different input shapes of the same length, but some have
+    # ones, broadcast to the correct shape.
+
+    for input_shapes in shapes:
+        assert_broadcast_correct(input_shapes)
+        assert_broadcast_correct(input_shapes[::-1])
+
+
+@pytest.mark.parametrize(
+    "shapes",
+    [
+        [[(), (3,)]],
+        [[(3,), (3, 3)]],
+        [[(3,), (3, 1)]],
+        [[(1,), (3, 3)]],
+        [[(), (3, 3)]],
+        [[(1, 1), (3,)]],
+        [[(1,), (3, 1)]],
+        [[(1,), (1, 3)]],
+        [[(), (1, 3)]],
+        [[(), (3, 1)]],
+        [[(), (0,)]],
+        [[(0,), (0, 0)]],
+        [[(0,), (0, 1)]],
+        [[(1,), (0, 0)]],
+        [[(), (0, 0)]],
+        [[(1, 1), (0,)]],
+        [[(1,), (0, 1)]],
+        [[(1,), (1, 0)]],
+        [[(), (1, 0)]],
+        [[(), (0, 1)]],
+    ],
+)
+def test_broadcast_arrays_different_len_shapes(shapes):
+    # Check that two different input shapes (of different lengths) broadcast
+    # to the correct shape.
+
+    for input_shapes in shapes:
+        assert_broadcast_correct(input_shapes)
+        assert_broadcast_correct(input_shapes[::-1])
+
+
+@pytest.mark.parametrize(
+    "shapes",
+    [
+        [[(3,), (4,)]],
+        [[(2, 3), (2,)]],
+        [[(3,), (3,), (4,)]],
+        [[(1, 3, 4), (2, 3, 3)]],
+    ],
+)
+def test_incompatible_shapes_raise_valueerror(shapes):
+    for input_shapes in shapes:
+        assert_broadcast_arrays_raise(input_shapes)
+        assert_broadcast_arrays_raise(input_shapes[::-1])