Skip to content

Commit 55361c7

Browse files
sklamsklam
authored
Merge pull request #7660 from guilhermeleobas/guilhermeleobas/np-broadcast_arrays
Add support for np.broadcast_arrays
2 parents 637a750 + a7292d1 commit 55361c7

3 files changed

Lines changed: 256 additions & 0 deletions

File tree

‎docs/source/reference/numpysupported.rst‎

Lines changed: 1 addition & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -435,6 +435,7 @@ The following top-level functions are supported:
435435
* :func:`numpy.bincount`
436436
* :func:`numpy.blackman`
437437
* :func:`numpy.broadcast_to` (only the 2 first arguments)
438+
* :func:`numpy.broadcast_arrays` (only the first argument)
438439
* :func:`numpy.column_stack`
439440
* :func:`numpy.concatenate`
440441
* :func:`numpy.convolve` (only the 2 first arguments)

‎numba/np/arrayobj.py‎

Lines changed: 73 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -1382,6 +1382,79 @@ def impl(array, shape):
13821382
return impl
13831383

13841384

1385+
@register_jitable
1386+
def numpy_broadcast_shapes_list(r, m, shape):
1387+
for i in range(len(shape)):
1388+
k = m - len(shape) + i
1389+
tmp = shape[i]
1390+
if tmp == 1:
1391+
continue
1392+
if r[k] == 1:
1393+
r[k] = tmp
1394+
elif r[k] != tmp:
1395+
raise ValueError("shape mismatch: objects"
1396+
" cannot be broadcast"
1397+
" to a single shape")
1398+
1399+
1400+
@overload(np.broadcast_arrays)
1401+
def numpy_broadcast_arrays(*args):
1402+
1403+
for idx, arg in enumerate(args):
1404+
if not type_can_asarray(arg):
1405+
raise errors.TypingError(f'Argument "{idx}" must '
1406+
'be array-like')
1407+
1408+
unified_dtype = None
1409+
dt = None
1410+
for arg in args:
1411+
if isinstance(arg, (types.Array, types.BaseTuple)):
1412+
dt = arg.dtype
1413+
else:
1414+
dt = arg
1415+
1416+
if unified_dtype is None:
1417+
unified_dtype = dt
1418+
elif unified_dtype != dt:
1419+
raise errors.TypingError('Mismatch of argument types. Numba cannot '
1420+
'broadcast arrays with different types. '
1421+
f'Got {args}')
1422+
1423+
# number of dimensions
1424+
m = 0
1425+
for idx, arg in enumerate(args):
1426+
if isinstance(arg, types.ArrayCompatible):
1427+
m = max(m, arg.ndim)
1428+
elif isinstance(arg, (types.Number, types.Boolean, types.BaseTuple)):
1429+
m = max(m, 1)
1430+
else:
1431+
raise errors.TypingError(f'Unhandled type {arg}')
1432+
1433+
tup_init = (0,) * m
1434+
1435+
def impl(*args):
1436+
# find out the output shape
1437+
# we can't call np.broadcast_shapes here since args may have arrays
1438+
# with different shapes and it is not possible to create a list
1439+
# with those shapes dynamically
1440+
shape = [1] * m
1441+
for array in literal_unroll(args):
1442+
numpy_broadcast_shapes_list(shape, m, np.asarray(array).shape)
1443+
1444+
tup = tup_init
1445+
1446+
for i in range(m):
1447+
tup = tuple_setitem(tup, i, shape[i])
1448+
1449+
# numpy checks if the input arrays have the same shape as `shape`
1450+
outs = []
1451+
for array in literal_unroll(args):
1452+
outs.append(np.broadcast_to(np.asarray(array), tup))
1453+
return outs
1454+
1455+
return impl
1456+
1457+
13851458
def fancy_setslice(context, builder, sig, args, index_types, indices):
13861459
"""
13871460
Implement slice assignment for arrays. This implementation works for

‎numba/tests/test_array_manipulation.py‎

Lines changed: 182 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -52,6 +52,10 @@ def numpy_broadcast_to(arr, shape):
5252
return np.broadcast_to(arr, shape)
5353

5454

55+
def numpy_broadcast_arrays(*args):
56+
return np.broadcast_arrays(*args)
57+
58+
5559
def numpy_broadcast_to_indexing(arr, shape, idx):
5660
return np.broadcast_to(arr, shape)[idx]
5761

@@ -965,6 +969,184 @@ def test_argwhere_array_like(self):
965969
got = cfunc(a)
966970
self.assertPreciseEqual(expected, got)
967971

972+
def broadcast_arrays_assert_correct_shape(self, input_shapes, expected_shape):
973+
# Broadcast a list of arrays with the given input shapes and check the
974+
# common output shape.
975+
pyfunc = numpy_broadcast_arrays
976+
cfunc = jit(nopython=True)(pyfunc)
977+
978+
inarrays = [np.zeros(s) for s in input_shapes]
979+
outarrays = cfunc(*inarrays)
980+
expected = [expected_shape] * len(inarrays)
981+
got = [a.shape for a in outarrays]
982+
self.assertPreciseEqual(expected, got)
983+
984+
def test_broadcast_arrays_same_input_shapes(self):
985+
# Tests taken from
986+
# https://github.com/numpy/numpy/blob/623bc1fae1d47df24e7f1e29321d0c0ba2771ce0/numpy/lib/tests/test_stride_tricks.py#L83-L107 # noqa: E501
987+
# Check that the final shape is just the input shape.
988+
pyfunc = numpy_broadcast_arrays
989+
cfunc = jit(nopython=True)(pyfunc)
990+
991+
data = [
992+
# (),
993+
(1,),
994+
(3,),
995+
(0, 1),
996+
(0, 3),
997+
(1, 0),
998+
(3, 0),
999+
(1, 3),
1000+
(3, 1),
1001+
(3, 3),
1002+
]
1003+
for shape in data:
1004+
input_shapes = [shape]
1005+
# Single input.
1006+
self.broadcast_arrays_assert_correct_shape(input_shapes, shape)
1007+
# Double input.
1008+
input_shapes2 = [shape, shape]
1009+
self.broadcast_arrays_assert_correct_shape(input_shapes2, shape)
1010+
# Triple input.
1011+
input_shapes3 = [shape, shape, shape]
1012+
self.broadcast_arrays_assert_correct_shape(input_shapes3, shape)
1013+
1014+
def test_broadcast_arrays_two_compatible_by_ones_input_shapes(self):
1015+
# Tests taken from
1016+
# https://github.com/numpy/numpy/blob/623bc1fae1d47df24e7f1e29321d0c0ba2771ce0/numpy/lib/tests/test_stride_tricks.py#L110-L132
1017+
# Check that two different input shapes of the same length, but some have
1018+
# ones, broadcast to the correct shape.
1019+
1020+
data = [
1021+
[[(1,), (3,)], (3,)],
1022+
[[(1, 3), (3, 3)], (3, 3)],
1023+
[[(3, 1), (3, 3)], (3, 3)],
1024+
[[(1, 3), (3, 1)], (3, 3)],
1025+
[[(1, 1), (3, 3)], (3, 3)],
1026+
[[(1, 1), (1, 3)], (1, 3)],
1027+
[[(1, 1), (3, 1)], (3, 1)],
1028+
[[(1, 0), (0, 0)], (0, 0)],
1029+
[[(0, 1), (0, 0)], (0, 0)],
1030+
[[(1, 0), (0, 1)], (0, 0)],
1031+
[[(1, 1), (0, 0)], (0, 0)],
1032+
[[(1, 1), (1, 0)], (1, 0)],
1033+
[[(1, 1), (0, 1)], (0, 1)],
1034+
]
1035+
for input_shapes, expected_shape in data:
1036+
self.broadcast_arrays_assert_correct_shape(input_shapes, expected_shape)
1037+
# Reverse the input shapes since broadcasting should be symmetric.
1038+
self.broadcast_arrays_assert_correct_shape(input_shapes[::-1], expected_shape)
1039+
1040+
def test_broadcast_arrays_two_compatible_by_prepending_ones_input_shapes(self):
1041+
# Tests taken from
1042+
# https://github.com/numpy/numpy/blob/623bc1fae1d47df24e7f1e29321d0c0ba2771ce0/numpy/lib/tests/test_stride_tricks.py#L135-L164
1043+
# Check that two different input shapes (of different lengths) broadcast
1044+
# to the correct shape.
1045+
1046+
data = [
1047+
[[(), (3,)], (3,)],
1048+
[[(3,), (3, 3)], (3, 3)],
1049+
[[(3,), (3, 1)], (3, 3)],
1050+
[[(1,), (3, 3)], (3, 3)],
1051+
[[(), (3, 3)], (3, 3)],
1052+
[[(1, 1), (3,)], (1, 3)],
1053+
[[(1,), (3, 1)], (3, 1)],
1054+
[[(1,), (1, 3)], (1, 3)],
1055+
[[(), (1, 3)], (1, 3)],
1056+
[[(), (3, 1)], (3, 1)],
1057+
[[(), (0,)], (0,)],
1058+
[[(0,), (0, 0)], (0, 0)],
1059+
[[(0,), (0, 1)], (0, 0)],
1060+
[[(1,), (0, 0)], (0, 0)],
1061+
[[(), (0, 0)], (0, 0)],
1062+
[[(1, 1), (0,)], (1, 0)],
1063+
[[(1,), (0, 1)], (0, 1)],
1064+
[[(1,), (1, 0)], (1, 0)],
1065+
[[(), (1, 0)], (1, 0)],
1066+
[[(), (0, 1)], (0, 1)],
1067+
]
1068+
for input_shapes, expected_shape in data:
1069+
self.broadcast_arrays_assert_correct_shape(input_shapes, expected_shape)
1070+
# Reverse the input shapes since broadcasting should be symmetric.
1071+
self.broadcast_arrays_assert_correct_shape(input_shapes[::-1], expected_shape)
1072+
1073+
def test_broadcast_arrays_scalar_input(self):
1074+
pyfunc = numpy_broadcast_arrays
1075+
cfunc = jit(nopython=True)(pyfunc)
1076+
data = [
1077+
[[True, False], (1,)],
1078+
[[1, 2], (1,)],
1079+
[[(1, 2), 2], (2,)],
1080+
]
1081+
for inarrays, expected_shape in data:
1082+
outarrays = cfunc(*inarrays)
1083+
got = [a.shape for a in outarrays]
1084+
expected = [expected_shape] * len(inarrays)
1085+
self.assertPreciseEqual(expected, got)
1086+
1087+
def test_broadcast_arrays_tuple_input(self):
1088+
pyfunc = numpy_broadcast_arrays
1089+
cfunc = jit(nopython=True)(pyfunc)
1090+
outarrays = cfunc((123, 456), (789,))
1091+
expected = [(2,), (2,)]
1092+
got = [a.shape for a in outarrays]
1093+
self.assertPreciseEqual(expected, got)
1094+
1095+
def test_broadcast_arrays_non_array_input(self):
1096+
pyfunc = numpy_broadcast_arrays
1097+
cfunc = jit(nopython=True)(pyfunc)
1098+
outarrays = cfunc(np.intp(2), np.zeros((1, 3), dtype=np.intp))
1099+
expected = [(1, 3), (1, 3)]
1100+
got = [a.shape for a in outarrays]
1101+
self.assertPreciseEqual(expected, got)
1102+
1103+
def test_broadcast_arrays_invalid_mixed_input_types(self):
1104+
pyfunc = numpy_broadcast_arrays
1105+
cfunc = jit(nopython=True)(pyfunc)
1106+
1107+
self.disable_leak_check()
1108+
1109+
with self.assertRaises(TypingError) as raises:
1110+
arr = np.arange(6).reshape((2, 3))
1111+
b = True
1112+
cfunc(arr, b)
1113+
self.assertIn('Mismatch of argument types', str(raises.exception))
1114+
1115+
def test_broadcast_arrays_invalid_input(self):
1116+
pyfunc = numpy_broadcast_arrays
1117+
cfunc = jit(nopython=True)(pyfunc)
1118+
1119+
self.disable_leak_check()
1120+
1121+
with self.assertRaises(TypingError) as raises:
1122+
arr = np.zeros(3, dtype=np.int64)
1123+
s = 'hello world'
1124+
cfunc(arr, s)
1125+
self.assertIn('Argument "1" must be array-like', str(raises.exception))
1126+
1127+
def test_broadcast_arrays_incompatible_shapes_raise_valueerror(self):
1128+
# Check that a ValueError is raised for incompatible shapes.
1129+
pyfunc = numpy_broadcast_arrays
1130+
cfunc = jit(nopython=True)(pyfunc)
1131+
1132+
self.disable_leak_check()
1133+
1134+
data = [
1135+
[(3,), (4,)],
1136+
[(2, 3), (2,)],
1137+
[(3,), (3,), (4,)],
1138+
[(1, 3, 4), (2, 3, 3)],
1139+
]
1140+
for input_shapes in data:
1141+
for shape in [input_shapes, input_shapes[::-1]]:
1142+
# Reverse the input shapes since broadcasting should be symmetric.
1143+
with self.assertRaises(ValueError) as raises:
1144+
inarrays = [np.zeros(s) for s in shape]
1145+
cfunc(*inarrays)
1146+
self.assertIn("shape mismatch: objects cannot be broadcast to a single shape",
1147+
str(raises.exception))
1148+
1149+
9681150

9691151
if __name__ == '__main__':
9701152
unittest.main()

0 commit comments

Comments
 (0)