sklam · Apr 1, 2022 · Dec 15, 2021 · Dec 15, 2021 · Dec 15, 2021 · Dec 16, 2021
diff --git a/docs/source/reference/numpysupported.rst b/docs/source/reference/numpysupported.rst
@@ -435,6 +435,7 @@ The following top-level functions are supported:
 * :func:`numpy.bincount`
 * :func:`numpy.blackman`
 * :func:`numpy.broadcast_to` (only the 2 first arguments)
+* :func:`numpy.broadcast_arrays` (only the first argument)
 * :func:`numpy.column_stack`
 * :func:`numpy.concatenate`
 * :func:`numpy.convolve` (only the 2 first arguments)

diff --git a/numba/np/arrayobj.py b/numba/np/arrayobj.py
@@ -1382,6 +1382,79 @@ def impl(array, shape):
     return impl
 
 
+@register_jitable
+def numpy_broadcast_shapes_list(r, m, shape):
+    for i in range(len(shape)):
+        k = m - len(shape) + i
+        tmp = shape[i]
+        if tmp == 1:
+            continue
+        if r[k] == 1:
+            r[k] = tmp
+        elif r[k] != tmp:
+            raise ValueError("shape mismatch: objects"
+                             " cannot be broadcast"
+                             " to a single shape")
+
+
+@overload(np.broadcast_arrays)
+def numpy_broadcast_arrays(*args):
+
+    for idx, arg in enumerate(args):
+        if not type_can_asarray(arg):
+            raise errors.TypingError(f'Argument "{idx}" must '
+                                     'be array-like')
+
+    unified_dtype = None
+    dt = None
+    for arg in args:
+        if isinstance(arg, (types.Array, types.BaseTuple)):
+            dt = arg.dtype
+        else:
+            dt = arg
+
+        if unified_dtype is None:
+            unified_dtype = dt
+        elif unified_dtype != dt:
+            raise errors.TypingError('Mismatch of argument types. Numba cannot '
+                                     'broadcast arrays with different types. '
+                                     f'Got {args}')
+
+    # number of dimensions
+    m = 0
+    for idx, arg in enumerate(args):
+        if isinstance(arg, types.ArrayCompatible):
+            m = max(m, arg.ndim)
+        elif isinstance(arg, (types.Number, types.Boolean, types.BaseTuple)):
+            m = max(m, 1)
+        else:
+            raise errors.TypingError(f'Unhandled type {arg}')
+
+    tup_init = (0,) * m
+
+    def impl(*args):
+        # find out the output shape
+        # we can't call np.broadcast_shapes here since args may have arrays
+        # with different shapes and it is not possible to create a list
+        # with those shapes dynamically
+        shape = [1] * m
+        for array in literal_unroll(args):
+            numpy_broadcast_shapes_list(shape, m, np.asarray(array).shape)
+
+        tup = tup_init
+
+        for i in range(m):
+            tup = tuple_setitem(tup, i, shape[i])
+
+        # numpy checks if the input arrays have the same shape as `shape`
+        outs = []
+        for array in literal_unroll(args):
+            outs.append(np.broadcast_to(np.asarray(array), tup))
+        return outs
+
+    return impl
+
+
 def fancy_setslice(context, builder, sig, args, index_types, indices):
     """
     Implement slice assignment for arrays.  This implementation works for

diff --git a/numba/tests/test_array_manipulation.py b/numba/tests/test_array_manipulation.py
@@ -52,6 +52,10 @@ def numpy_broadcast_to(arr, shape):
     return np.broadcast_to(arr, shape)
 
 
+def numpy_broadcast_arrays(*args):
+    return np.broadcast_arrays(*args)
+
+
 def numpy_broadcast_to_indexing(arr, shape, idx):
     return np.broadcast_to(arr, shape)[idx]
 
@@ -965,6 +969,184 @@ def test_argwhere_array_like(self):
             got = cfunc(a)
             self.assertPreciseEqual(expected, got)
 
+    def broadcast_arrays_assert_correct_shape(self, input_shapes, expected_shape):
+        # Broadcast a list of arrays with the given input shapes and check the
+        # common output shape.
+        pyfunc = numpy_broadcast_arrays
+        cfunc = jit(nopython=True)(pyfunc)
+
+        inarrays = [np.zeros(s) for s in input_shapes]
+        outarrays = cfunc(*inarrays)
+        expected = [expected_shape] * len(inarrays)
+        got = [a.shape for a in outarrays]
+        self.assertPreciseEqual(expected, got)
+
+    def test_broadcast_arrays_same_input_shapes(self):
+        # Tests taken from
+        # https://github.com/numpy/numpy/blob/623bc1fae1d47df24e7f1e29321d0c0ba2771ce0/numpy/lib/tests/test_stride_tricks.py#L83-L107  # noqa: E501
+        # Check that the final shape is just the input shape.
+        pyfunc = numpy_broadcast_arrays
+        cfunc = jit(nopython=True)(pyfunc)
+
+        data = [
+            # (),
+            (1,),
+            (3,),
+            (0, 1),
+            (0, 3),
+            (1, 0),
+            (3, 0),
+            (1, 3),
+            (3, 1),
+            (3, 3),
+        ]
+        for shape in data:
+            input_shapes = [shape]
+            # Single input.
+            self.broadcast_arrays_assert_correct_shape(input_shapes, shape)
+            # Double input.
+            input_shapes2 = [shape, shape]
+            self.broadcast_arrays_assert_correct_shape(input_shapes2, shape)
+            # Triple input.
+            input_shapes3 = [shape, shape, shape]
+            self.broadcast_arrays_assert_correct_shape(input_shapes3, shape)
+
+    def test_broadcast_arrays_two_compatible_by_ones_input_shapes(self):
+        # Tests taken from
+        # https://github.com/numpy/numpy/blob/623bc1fae1d47df24e7f1e29321d0c0ba2771ce0/numpy/lib/tests/test_stride_tricks.py#L110-L132
+        # Check that two different input shapes of the same length, but some have
+        # ones, broadcast to the correct shape.
+
+        data = [
+            [[(1,), (3,)], (3,)],
+            [[(1, 3), (3, 3)], (3, 3)],
+            [[(3, 1), (3, 3)], (3, 3)],
+            [[(1, 3), (3, 1)], (3, 3)],
+            [[(1, 1), (3, 3)], (3, 3)],
+            [[(1, 1), (1, 3)], (1, 3)],
+            [[(1, 1), (3, 1)], (3, 1)],
+            [[(1, 0), (0, 0)], (0, 0)],
+            [[(0, 1), (0, 0)], (0, 0)],
+            [[(1, 0), (0, 1)], (0, 0)],
+            [[(1, 1), (0, 0)], (0, 0)],
+            [[(1, 1), (1, 0)], (1, 0)],
+            [[(1, 1), (0, 1)], (0, 1)],
+        ]
+        for input_shapes, expected_shape in data:
+            self.broadcast_arrays_assert_correct_shape(input_shapes, expected_shape)
+            # Reverse the input shapes since broadcasting should be symmetric.
+            self.broadcast_arrays_assert_correct_shape(input_shapes[::-1], expected_shape)
+
+    def test_broadcast_arrays_two_compatible_by_prepending_ones_input_shapes(self):
+        # Tests taken from
+        # https://github.com/numpy/numpy/blob/623bc1fae1d47df24e7f1e29321d0c0ba2771ce0/numpy/lib/tests/test_stride_tricks.py#L135-L164
+        # Check that two different input shapes (of different lengths) broadcast
+        # to the correct shape.
+
+        data = [
+            [[(), (3,)], (3,)],
+            [[(3,), (3, 3)], (3, 3)],
+            [[(3,), (3, 1)], (3, 3)],
+            [[(1,), (3, 3)], (3, 3)],
+            [[(), (3, 3)], (3, 3)],
+            [[(1, 1), (3,)], (1, 3)],
+            [[(1,), (3, 1)], (3, 1)],
+            [[(1,), (1, 3)], (1, 3)],
+            [[(), (1, 3)], (1, 3)],
+            [[(), (3, 1)], (3, 1)],
+            [[(), (0,)], (0,)],
+            [[(0,), (0, 0)], (0, 0)],
+            [[(0,), (0, 1)], (0, 0)],
+            [[(1,), (0, 0)], (0, 0)],
+            [[(), (0, 0)], (0, 0)],
+            [[(1, 1), (0,)], (1, 0)],
+            [[(1,), (0, 1)], (0, 1)],
+            [[(1,), (1, 0)], (1, 0)],
+            [[(), (1, 0)], (1, 0)],
+            [[(), (0, 1)], (0, 1)],
+        ]
+        for input_shapes, expected_shape in data:
+            self.broadcast_arrays_assert_correct_shape(input_shapes, expected_shape)
+            # Reverse the input shapes since broadcasting should be symmetric.
+            self.broadcast_arrays_assert_correct_shape(input_shapes[::-1], expected_shape)
+
+    def test_broadcast_arrays_scalar_input(self):
+        pyfunc = numpy_broadcast_arrays
+        cfunc = jit(nopython=True)(pyfunc)
+        data = [
+            [[True, False], (1,)],
+            [[1, 2], (1,)],
+            [[(1, 2), 2], (2,)],
+        ]
+        for inarrays, expected_shape in data:
+            outarrays = cfunc(*inarrays)
+            got = [a.shape for a in outarrays]
+            expected = [expected_shape] * len(inarrays)
+            self.assertPreciseEqual(expected, got)
+
+    def test_broadcast_arrays_tuple_input(self):
+        pyfunc = numpy_broadcast_arrays
+        cfunc = jit(nopython=True)(pyfunc)
+        outarrays = cfunc((123, 456), (789,))
+        expected = [(2,), (2,)]
+        got = [a.shape for a in outarrays]
+        self.assertPreciseEqual(expected, got)
+
+    def test_broadcast_arrays_non_array_input(self):
+        pyfunc = numpy_broadcast_arrays
+        cfunc = jit(nopython=True)(pyfunc)
+        outarrays = cfunc(np.intp(2), np.zeros((1, 3), dtype=np.intp))
+        expected = [(1, 3), (1, 3)]
+        got = [a.shape for a in outarrays]
+        self.assertPreciseEqual(expected, got)
+
+    def test_broadcast_arrays_invalid_mixed_input_types(self):
+        pyfunc = numpy_broadcast_arrays
+        cfunc = jit(nopython=True)(pyfunc)
+
+        self.disable_leak_check()
+
+        with self.assertRaises(TypingError) as raises:
+            arr = np.arange(6).reshape((2, 3))
+            b = True
+            cfunc(arr, b)
+        self.assertIn('Mismatch of argument types', str(raises.exception))
+
+    def test_broadcast_arrays_invalid_input(self):
+        pyfunc = numpy_broadcast_arrays
+        cfunc = jit(nopython=True)(pyfunc)
+
+        self.disable_leak_check()
+
+        with self.assertRaises(TypingError) as raises:
+            arr = np.zeros(3, dtype=np.int64)
+            s = 'hello world'
+            cfunc(arr, s)
+        self.assertIn('Argument "1" must be array-like', str(raises.exception))
+
+    def test_broadcast_arrays_incompatible_shapes_raise_valueerror(self):
+        # Check that a ValueError is raised for incompatible shapes.
+        pyfunc = numpy_broadcast_arrays
+        cfunc = jit(nopython=True)(pyfunc)
+
+        self.disable_leak_check()
+
+        data = [
+            [(3,), (4,)],
+            [(2, 3), (2,)],
+            [(3,), (3,), (4,)],
+            [(1, 3, 4), (2, 3, 3)],
+        ]
+        for input_shapes in data:
+            for shape in [input_shapes, input_shapes[::-1]]:
+                # Reverse the input shapes since broadcasting should be symmetric.
+                with self.assertRaises(ValueError) as raises:
+                    inarrays = [np.zeros(s) for s in shape]
+                    cfunc(*inarrays)
+                self.assertIn("shape mismatch: objects cannot be broadcast to a single shape",
+                              str(raises.exception))
+
+
 
 if __name__ == '__main__':
     unittest.main()