Skip to content

Commit

Permalink
Merge pull request #1594 from IntelPython/broadcast_arrays
Browse files Browse the repository at this point in the history
Implemented dpnp.broadcast_arrays function.
  • Loading branch information
npolina4 authored Oct 17, 2023
2 parents c1bf9ac + 3432716 commit b469ed7
Show file tree
Hide file tree
Showing 5 changed files with 220 additions and 48 deletions.
69 changes: 61 additions & 8 deletions dpnp/dpnp_iface_manipulation.py
Original file line number Diff line number Diff line change
Expand Up @@ -54,6 +54,7 @@
"atleast_1d",
"atleast_2d",
"atleast_3d",
"broadcast_arrays",
"broadcast_to",
"concatenate",
"copyto",
Expand Down Expand Up @@ -309,6 +310,55 @@ def atleast_3d(*arys):
return res


def broadcast_arrays(*args, subok=False):
"""
Broadcast any number of arrays against each other.
For full documentation refer to :obj:`numpy.broadcast_arrays`.
Returns
-------
broadcasted : list of dpnp.ndarray
These arrays are views on the original arrays.
Limitations
-----------
Parameter `args` is supported as either :class:`dpnp.ndarray`
or :class:`dpctl.tensor.usm_ndarray`.
Otherwise ``TypeError`` exception will be raised.
Parameter `subok` is supported with default value.
Otherwise ``NotImplementedError`` exception will be raised.
See Also
--------
:obj:`dpnp.broadcast_to` : Broadcast an array to a new shape.
Examples
--------
>>> import dpnp as np
>>> x = np.array([[1, 2, 3]])
>>> y = np.array([[4], [5]])
>>> np.broadcast_arrays(x, y)
[array([[1, 2, 3],
[1, 2, 3]]), array([[4, 4, 4],
[5, 5, 5]])]
"""

if subok is not False:
raise NotImplementedError(f"subok={subok} is currently not supported")

if len(args) == 0:
return []

dpt_arrays = dpt.broadcast_arrays(
*[dpnp.get_usm_ndarray(array) for array in args]
)
return [
dpnp_array._create_from_usm_ndarray(usm_arr) for usm_arr in dpt_arrays
]


def broadcast_to(array, /, shape, subok=False):
"""
Broadcast an array to a new shape.
Expand All @@ -324,10 +374,15 @@ def broadcast_to(array, /, shape, subok=False):
-----------
Parameter `array` is supported as either :class:`dpnp.ndarray`
or :class:`dpctl.tensor.usm_ndarray`.
Otherwise ``TypeError`` exception will be raised.
Parameter `subok` is supported with default value.
Otherwise the function will be executed sequentially on CPU.
Otherwise ``NotImplementedError`` exception will be raised.
Input array data types of `array` is limited by supported DPNP :ref:`Data types`.
See Also
--------
:obj:`dpnp.broadcast_arrays` : Broadcast any number of arrays against each other.
Examples
--------
>>> import dpnp as dp
Expand All @@ -340,13 +395,11 @@ def broadcast_to(array, /, shape, subok=False):
"""

if subok is not False:
pass
elif dpnp.is_supported_array_type(array):
dpt_array = dpnp.get_usm_ndarray(array)
new_array = dpt.broadcast_to(dpt_array, shape)
return dpnp_array._create_from_usm_ndarray(new_array)
raise NotImplementedError(f"subok={subok} is currently not supported")

return call_origin(numpy.broadcast_to, array, shape=shape, subok=subok)
dpt_array = dpnp.get_usm_ndarray(array)
new_array = dpt.broadcast_to(dpt_array, shape)
return dpnp_array._create_from_usm_ndarray(new_array)


def concatenate(
Expand All @@ -367,7 +420,7 @@ def concatenate(
Each array in `arrays` is supported as either :class:`dpnp.ndarray`
or :class:`dpctl.tensor.usm_ndarray`. Otherwise ``TypeError`` exception
will be raised.
Parameters `out` and `dtype are supported with default value.
Parameters `out` and `dtype` are supported with default value.
Otherwise the function will be executed sequentially on CPU.
See Also
Expand Down
16 changes: 0 additions & 16 deletions tests/skipped_tests.tbl
Original file line number Diff line number Diff line change
Expand Up @@ -119,7 +119,6 @@ tests/third_party/cupy/core_tests/test_ndarray_conversion.py::TestNdarrayToBytes
tests/third_party/cupy/core_tests/test_ndarray_copy_and_view.py::TestArrayCopyAndView::test_astype
tests/third_party/cupy/core_tests/test_ndarray_copy_and_view.py::TestArrayCopyAndView::test_astype_type
tests/third_party/cupy/core_tests/test_ndarray_copy_and_view.py::TestArrayCopyAndView::test_astype_strides
tests/third_party/cupy/core_tests/test_ndarray_copy_and_view.py::TestArrayCopyAndView::test_astype_strides_broadcast
tests/third_party/cupy/core_tests/test_ndarray_copy_and_view.py::TestArrayCopyAndView::test_astype_strides_negative
tests/third_party/cupy/core_tests/test_ndarray_copy_and_view.py::TestArrayCopyAndView::test_astype_strides_swapped
tests/third_party/cupy/core_tests/test_ndarray_copy_and_view.py::TestArrayCopyAndView::test_astype_type_c_contiguous_no_copy
Expand Down Expand Up @@ -427,35 +426,20 @@ tests/third_party/cupy/logic_tests/test_comparison.py::TestArrayEqual::test_arra
tests/third_party/cupy/logic_tests/test_comparison.py::TestArrayEqual::test_array_equal_not_equal

tests/third_party/cupy/manipulation_tests/test_dims.py::TestBroadcast_param_0_{shapes=[(), ()]}::test_broadcast
tests/third_party/cupy/manipulation_tests/test_dims.py::TestBroadcast_param_0_{shapes=[(), ()]}::test_broadcast_arrays
tests/third_party/cupy/manipulation_tests/test_dims.py::TestBroadcast_param_10_{shapes=[(0, 1, 1, 0, 3), (5, 2, 0, 1, 0, 0, 3), (2, 1, 0, 0, 0, 3)]}::test_broadcast
tests/third_party/cupy/manipulation_tests/test_dims.py::TestBroadcast_param_10_{shapes=[(0, 1, 1, 0, 3), (5, 2, 0, 1, 0, 0, 3), (2, 1, 0, 0, 0, 3)]}::test_broadcast_arrays
tests/third_party/cupy/manipulation_tests/test_dims.py::TestBroadcast_param_1_{shapes=[(0,), (0,)]}::test_broadcast
tests/third_party/cupy/manipulation_tests/test_dims.py::TestBroadcast_param_1_{shapes=[(0,), (0,)]}::test_broadcast_arrays
tests/third_party/cupy/manipulation_tests/test_dims.py::TestBroadcast_param_2_{shapes=[(1,), (1,)]}::test_broadcast
tests/third_party/cupy/manipulation_tests/test_dims.py::TestBroadcast_param_2_{shapes=[(1,), (1,)]}::test_broadcast_arrays
tests/third_party/cupy/manipulation_tests/test_dims.py::TestBroadcast_param_3_{shapes=[(2,), (2,)]}::test_broadcast
tests/third_party/cupy/manipulation_tests/test_dims.py::TestBroadcast_param_3_{shapes=[(2,), (2,)]}::test_broadcast_arrays
tests/third_party/cupy/manipulation_tests/test_dims.py::TestBroadcast_param_4_{shapes=[(0,), (1,)]}::test_broadcast
tests/third_party/cupy/manipulation_tests/test_dims.py::TestBroadcast_param_4_{shapes=[(0,), (1,)]}::test_broadcast_arrays
tests/third_party/cupy/manipulation_tests/test_dims.py::TestBroadcast_param_5_{shapes=[(2, 3), (1, 3)]}::test_broadcast
tests/third_party/cupy/manipulation_tests/test_dims.py::TestBroadcast_param_5_{shapes=[(2, 3), (1, 3)]}::test_broadcast_arrays
tests/third_party/cupy/manipulation_tests/test_dims.py::TestBroadcast_param_6_{shapes=[(2, 1, 3, 4), (3, 1, 4)]}::test_broadcast
tests/third_party/cupy/manipulation_tests/test_dims.py::TestBroadcast_param_6_{shapes=[(2, 1, 3, 4), (3, 1, 4)]}::test_broadcast_arrays
tests/third_party/cupy/manipulation_tests/test_dims.py::TestBroadcast_param_7_{shapes=[(4, 3, 2, 3), (2, 3)]}::test_broadcast
tests/third_party/cupy/manipulation_tests/test_dims.py::TestBroadcast_param_7_{shapes=[(4, 3, 2, 3), (2, 3)]}::test_broadcast_arrays
tests/third_party/cupy/manipulation_tests/test_dims.py::TestBroadcast_param_8_{shapes=[(2, 0, 1, 1, 3), (2, 1, 0, 0, 3)]}::test_broadcast
tests/third_party/cupy/manipulation_tests/test_dims.py::TestBroadcast_param_8_{shapes=[(2, 0, 1, 1, 3), (2, 1, 0, 0, 3)]}::test_broadcast_arrays
tests/third_party/cupy/manipulation_tests/test_dims.py::TestBroadcast_param_9_{shapes=[(0, 1, 1, 3), (2, 1, 0, 0, 3)]}::test_broadcast
tests/third_party/cupy/manipulation_tests/test_dims.py::TestBroadcast_param_9_{shapes=[(0, 1, 1, 3), (2, 1, 0, 0, 3)]}::test_broadcast_arrays
tests/third_party/cupy/manipulation_tests/test_dims.py::TestInvalidBroadcast_param_0_{shapes=[(3,), (2,)]}::test_invalid_broadcast
tests/third_party/cupy/manipulation_tests/test_dims.py::TestInvalidBroadcast_param_0_{shapes=[(3,), (2,)]}::test_invalid_broadcast_arrays
tests/third_party/cupy/manipulation_tests/test_dims.py::TestInvalidBroadcast_param_1_{shapes=[(3, 2), (2, 3)]}::test_invalid_broadcast
tests/third_party/cupy/manipulation_tests/test_dims.py::TestInvalidBroadcast_param_1_{shapes=[(3, 2), (2, 3)]}::test_invalid_broadcast_arrays
tests/third_party/cupy/manipulation_tests/test_dims.py::TestInvalidBroadcast_param_2_{shapes=[(3, 2), (3, 4)]}::test_invalid_broadcast
tests/third_party/cupy/manipulation_tests/test_dims.py::TestInvalidBroadcast_param_2_{shapes=[(3, 2), (3, 4)]}::test_invalid_broadcast_arrays
tests/third_party/cupy/manipulation_tests/test_dims.py::TestInvalidBroadcast_param_3_{shapes=[(0,), (2,)]}::test_invalid_broadcast
tests/third_party/cupy/manipulation_tests/test_dims.py::TestInvalidBroadcast_param_3_{shapes=[(0,), (2,)]}::test_invalid_broadcast_arrays

tests/third_party/cupy/manipulation_tests/test_shape.py::TestRavel::test_ravel2
tests/third_party/cupy/manipulation_tests/test_shape.py::TestRavel::test_ravel3
Expand Down
18 changes: 1 addition & 17 deletions tests/skipped_tests_gpu.tbl
Original file line number Diff line number Diff line change
Expand Up @@ -248,7 +248,6 @@ tests/third_party/cupy/core_tests/test_ndarray_complex_ops.py::TestScalarConvers
tests/third_party/cupy/core_tests/test_ndarray_copy_and_view.py::TestArrayCopyAndView::test_astype
tests/third_party/cupy/core_tests/test_ndarray_copy_and_view.py::TestArrayCopyAndView::test_astype_type
tests/third_party/cupy/core_tests/test_ndarray_copy_and_view.py::TestArrayCopyAndView::test_astype_strides
tests/third_party/cupy/core_tests/test_ndarray_copy_and_view.py::TestArrayCopyAndView::test_astype_strides_broadcast
tests/third_party/cupy/core_tests/test_ndarray_copy_and_view.py::TestArrayCopyAndView::test_astype_strides_negative
tests/third_party/cupy/core_tests/test_ndarray_copy_and_view.py::TestArrayCopyAndView::test_astype_strides_swapped
tests/third_party/cupy/core_tests/test_ndarray_copy_and_view.py::TestArrayCopyAndView::test_astype_type_c_contiguous_no_copy
Expand Down Expand Up @@ -572,36 +571,21 @@ tests/third_party/cupy/logic_tests/test_comparison.py::TestArrayEqual::test_arra
tests/third_party/cupy/logic_tests/test_comparison.py::TestArrayEqual::test_array_equal_not_equal

tests/third_party/cupy/manipulation_tests/test_dims.py::TestBroadcast_param_0_{shapes=[(), ()]}::test_broadcast
tests/third_party/cupy/manipulation_tests/test_dims.py::TestBroadcast_param_0_{shapes=[(), ()]}::test_broadcast_arrays
tests/third_party/cupy/manipulation_tests/test_dims.py::TestBroadcast_param_10_{shapes=[(0, 1, 1, 0, 3), (5, 2, 0, 1, 0, 0, 3), (2, 1, 0, 0, 0, 3)]}::test_broadcast
tests/third_party/cupy/manipulation_tests/test_dims.py::TestBroadcast_param_10_{shapes=[(0, 1, 1, 0, 3), (5, 2, 0, 1, 0, 0, 3), (2, 1, 0, 0, 0, 3)]}::test_broadcast_arrays
tests/third_party/cupy/manipulation_tests/test_dims.py::TestBroadcast_param_1_{shapes=[(0,), (0,)]}::test_broadcast
tests/third_party/cupy/manipulation_tests/test_dims.py::TestBroadcast_param_1_{shapes=[(0,), (0,)]}::test_broadcast_arrays
tests/third_party/cupy/manipulation_tests/test_dims.py::TestBroadcast_param_2_{shapes=[(1,), (1,)]}::test_broadcast
tests/third_party/cupy/manipulation_tests/test_dims.py::TestBroadcast_param_2_{shapes=[(1,), (1,)]}::test_broadcast_arrays
tests/third_party/cupy/manipulation_tests/test_dims.py::TestBroadcast_param_3_{shapes=[(2,), (2,)]}::test_broadcast
tests/third_party/cupy/manipulation_tests/test_dims.py::TestBroadcast_param_3_{shapes=[(2,), (2,)]}::test_broadcast_arrays
tests/third_party/cupy/manipulation_tests/test_dims.py::TestBroadcast_param_4_{shapes=[(0,), (1,)]}::test_broadcast
tests/third_party/cupy/manipulation_tests/test_dims.py::TestBroadcast_param_4_{shapes=[(0,), (1,)]}::test_broadcast_arrays
tests/third_party/cupy/manipulation_tests/test_dims.py::TestBroadcast_param_5_{shapes=[(2, 3), (1, 3)]}::test_broadcast
tests/third_party/cupy/manipulation_tests/test_dims.py::TestBroadcast_param_5_{shapes=[(2, 3), (1, 3)]}::test_broadcast_arrays
tests/third_party/cupy/manipulation_tests/test_dims.py::TestBroadcast_param_6_{shapes=[(2, 1, 3, 4), (3, 1, 4)]}::test_broadcast
tests/third_party/cupy/manipulation_tests/test_dims.py::TestBroadcast_param_6_{shapes=[(2, 1, 3, 4), (3, 1, 4)]}::test_broadcast_arrays
tests/third_party/cupy/manipulation_tests/test_dims.py::TestBroadcast_param_7_{shapes=[(4, 3, 2, 3), (2, 3)]}::test_broadcast
tests/third_party/cupy/manipulation_tests/test_dims.py::TestBroadcast_param_7_{shapes=[(4, 3, 2, 3), (2, 3)]}::test_broadcast_arrays
tests/third_party/cupy/manipulation_tests/test_dims.py::TestBroadcast_param_8_{shapes=[(2, 0, 1, 1, 3), (2, 1, 0, 0, 3)]}::test_broadcast
tests/third_party/cupy/manipulation_tests/test_dims.py::TestBroadcast_param_8_{shapes=[(2, 0, 1, 1, 3), (2, 1, 0, 0, 3)]}::test_broadcast_arrays
tests/third_party/cupy/manipulation_tests/test_dims.py::TestBroadcast_param_9_{shapes=[(0, 1, 1, 3), (2, 1, 0, 0, 3)]}::test_broadcast
tests/third_party/cupy/manipulation_tests/test_dims.py::TestBroadcast_param_9_{shapes=[(0, 1, 1, 3), (2, 1, 0, 0, 3)]}::test_broadcast_arrays

tests/third_party/cupy/manipulation_tests/test_dims.py::TestInvalidBroadcast_param_0_{shapes=[(3,), (2,)]}::test_invalid_broadcast
tests/third_party/cupy/manipulation_tests/test_dims.py::TestInvalidBroadcast_param_0_{shapes=[(3,), (2,)]}::test_invalid_broadcast_arrays
tests/third_party/cupy/manipulation_tests/test_dims.py::TestInvalidBroadcast_param_1_{shapes=[(3, 2), (2, 3)]}::test_invalid_broadcast
tests/third_party/cupy/manipulation_tests/test_dims.py::TestInvalidBroadcast_param_1_{shapes=[(3, 2), (2, 3)]}::test_invalid_broadcast_arrays
tests/third_party/cupy/manipulation_tests/test_dims.py::TestInvalidBroadcast_param_2_{shapes=[(3, 2), (3, 4)]}::test_invalid_broadcast
tests/third_party/cupy/manipulation_tests/test_dims.py::TestInvalidBroadcast_param_2_{shapes=[(3, 2), (3, 4)]}::test_invalid_broadcast_arrays
tests/third_party/cupy/manipulation_tests/test_dims.py::TestInvalidBroadcast_param_3_{shapes=[(0,), (2,)]}::test_invalid_broadcast
tests/third_party/cupy/manipulation_tests/test_dims.py::TestInvalidBroadcast_param_3_{shapes=[(0,), (2,)]}::test_invalid_broadcast_arrays

tests/third_party/cupy/manipulation_tests/test_shape.py::TestRavel::test_ravel2
tests/third_party/cupy/manipulation_tests/test_shape.py::TestRavel::test_ravel3
tests/third_party/cupy/manipulation_tests/test_shape.py::TestRavel::test_external_ravel
Expand Down
150 changes: 150 additions & 0 deletions tests/test_arraymanipulation.py
Original file line number Diff line number Diff line change
Expand Up @@ -778,3 +778,153 @@ def test_3D_array(self):
res = [dpnp.atleast_3d(a), dpnp.atleast_3d(b)]
desired = [a, b]
assert_array_equal(res, desired)


def assert_broadcast_correct(input_shapes):
np_arrays = [numpy.zeros(s, dtype="i1") for s in input_shapes]
out_np_arrays = numpy.broadcast_arrays(*np_arrays)
dpnp_arrays = [dpnp.asarray(Xnp) for Xnp in np_arrays]
out_dpnp_arrays = dpnp.broadcast_arrays(*dpnp_arrays)
for Xnp, X in zip(out_np_arrays, out_dpnp_arrays):
assert_array_equal(
Xnp, dpnp.asnumpy(X), err_msg=f"Failed for {input_shapes})"
)


def assert_broadcast_arrays_raise(input_shapes):
dpnp_arrays = [dpnp.asarray(numpy.zeros(s)) for s in input_shapes]
pytest.raises(ValueError, dpnp.broadcast_arrays, *dpnp_arrays)


def test_broadcast_arrays_same():
Xnp = numpy.arange(10)
Ynp = numpy.arange(10)
res_Xnp, res_Ynp = numpy.broadcast_arrays(Xnp, Ynp)
X = dpnp.asarray(Xnp)
Y = dpnp.asarray(Ynp)
res_X, res_Y = dpnp.broadcast_arrays(X, Y)
assert_array_equal(res_Xnp, dpnp.asnumpy(res_X))
assert_array_equal(res_Ynp, dpnp.asnumpy(res_Y))


def test_broadcast_arrays_one_off():
Xnp = numpy.array([[1, 2, 3]])
Ynp = numpy.array([[1], [2], [3]])
res_Xnp, res_Ynp = numpy.broadcast_arrays(Xnp, Ynp)
X = dpnp.asarray(Xnp)
Y = dpnp.asarray(Ynp)
res_X, res_Y = dpnp.broadcast_arrays(X, Y)
assert_array_equal(res_Xnp, dpnp.asnumpy(res_X))
assert_array_equal(res_Ynp, dpnp.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])


def test_broadcast_arrays_empty_input():
assert dpnp.broadcast_arrays() == []


def test_subok_error():
x = dpnp.ones((4))
with pytest.raises(NotImplementedError):
dpnp.broadcast_arrays(x, subok=True)
dpnp.broadcast_to(x, (4, 4), subok=True)
15 changes: 8 additions & 7 deletions tests/third_party/cupy/core_tests/test_ndarray_copy_and_view.py
Original file line number Diff line number Diff line change
Expand Up @@ -157,13 +157,14 @@ def test_astype_strides_swapped(self, xp, src_dtype, dst_dtype):
@testing.for_all_dtypes_combination(("src_dtype", "dst_dtype"))
@testing.numpy_cupy_array_equal()
def test_astype_strides_broadcast(self, xp, src_dtype, dst_dtype):
src, _ = xp.broadcast_arrays(
xp.empty((2,), dtype=src_dtype),
xp.empty((2, 3, 2), dtype=src_dtype),
)
return numpy.array(
astype_without_warning(src, dst_dtype, order="K").strides
)
src1 = testing.shaped_arange((2, 3, 2), xp, dtype=src_dtype)
src2 = testing.shaped_arange((2,), xp, dtype=src_dtype)
src, _ = xp.broadcast_arrays(src1, src2)
dst = astype_without_warning(src, dst_dtype, order="K")
strides = dst.strides
if xp is numpy:
strides = tuple(x // dst.itemsize for x in strides)
return strides

@pytest.mark.usefixtures("allow_fall_back_on_numpy")
@testing.for_all_dtypes()
Expand Down

0 comments on commit b469ed7

Please sign in to comment.