Checking if two arrays are broadcastable in python

Question

According to the documentation (https://docs.scipy.org/doc/numpy/reference/ufuncs.html), two arrays are broadcastable if one of the following conditions are true:

- The arrays all have exactly the same shape.
- The arrays all have the same number of dimensions and the length of each dimensions is either a common length or 1.
- Arrays that have too few dimensions can have their shapes prepended with a dimension of length 1 to satisfy property 2.

I tried to implement this in python, but having a trouble in understanding the 2nd and 3rd rule. Not getting the answer that I am expecting, but I want to know what error that I am making in my code and a possible solution to this error.

# Broadcastable if:
# The arrays all have exactly the same shape.

if a.shape == b.shape:
    result = True

# The arrays all have the same number of dimensions 

elif len(a.shape) == len(b.shape):

    # and the length of each dimensions is either a common length or 1.

    for i in len(a.shape):
        if len(a.shape[i] != 1):
            result = False
        else:
            result = True

# Arrays that have too few dimensions can have their shapes prepended with a dimension of length 1 to satisfy property 2

elif a.shape == () or b.shape == ():
    result = True


else:
    result = False

return result

Use `numpy.broadcast_arrays` and check if shapes are equal? – Divakar Nov 11 '17 at 22:35 — Divakar, Nov 11 '17 at 22:35

Warren Weckesser · Accepted Answer · 2017-11-12T04:14:16.563

10

Here's a concise expression to check if two arrays are broadcastable:

In [101]: import numpy as np

In [102]: a = np.zeros((3, 1, 5))

In [103]: b = np.zeros((4, 5))

In [104]: all((m == n) or (m == 1) or (n == 1) for m, n in zip(a.shape[::-1], b.shape[::-1]))
Out[104]: True

In [105]: b = np.zeros((5, 3))

In [106]: all((m == n) or (m == 1) or (n == 1) for m, n in zip(a.shape[::-1], b.shape[::-1]))
Out[106]: False

edited Nov 12 '17 at 04:14

answered Nov 12 '17 at 00:30

Warren Weckesser

110,654
19
194
214

Nice use of `zip_longest`. – hpaulj Nov 12 '17 at 02:24
1

Wouldn't `zip` be enough? – Paul Panzer Nov 12 '17 at 03:46
@PaulPanzer Yes, it would. `zip_longest` is used to pad the "shorter" shape with 1s, but those 1s will never cause a `False` in the sequence of comparisons, so they are not necessary. I've updated the code to use `zip`. Thanks for the suggestion! – Warren Weckesser Nov 12 '17 at 04:07
gives me a false negative with shapes (30, 1, 2) and (3,2). These arrays are broadcastable, yet this snippet returns false. – Adrian Nov 09 '20 at 18:20
@Adrian, I just tried it with `a = np.zeros((30, 1, 2)) ` and `b = np.zeros((3, 2))`, and the snippet returned True. – Warren Weckesser Nov 09 '20 at 19:11
Oh wait didn't see the `[:, :, - 1]` part! My bad. – Adrian Nov 09 '20 at 19:36
Perhaps it should be added that this is available in numpy after 1.20.0, simply use `np.broadcast_shapes` – nickpapior Apr 22 '22 at 06:25

hpaulj · Answer 2 · 2017-11-12T00:39:09.587

I haven't tried to implement this as code, but the way I reason things is:

compare the shapes
if one has fewer dimensions, add 1s at the start to match
again compare; replace all 1s with a corresponding dimension from the other array

For example, if I have a (3,1) and (3,) array. Expand the (3,) to (1,3). Now change both to (3,3).

In [180]: A = np.ones((3,1),int); B = np.arange(3)
In [181]: A.shape
Out[181]: (3, 1)
In [182]: B.shape
Out[182]: (3,)
In [183]: (A+B)
Out[183]: 
array([[1, 2, 3],
       [1, 2, 3],
       [1, 2, 3]])

In [184]: np.broadcast_arrays(A,B)
Out[184]: 
[array([[1, 1, 1],
        [1, 1, 1],
        [1, 1, 1]]), 
 array([[0, 1, 2],
        [0, 1, 2],
        [0, 1, 2]])]

In another recent question I use broadcast_to. I start with a n-d array, reduce one dimension to 1 with np.mean (and keepdims), and then re-expand it with broadcast_to.

Arbitrary N-dimensional repeat of N-dimensional numpy array

In a recent answer, https://stackoverflow.com/a/47243071/901925

a (2,3) is broadcasted with a (4,3) with

A[:, None] makes the (4,3) into a (4,1,3) (there's an implied final ':'). B is automatically expanded from (2,3) to (1,2,3). Now both can broadcast to (4,2,3).

So the (2,3) can't broadcast with (4,3) (same n-d but different values). But it can broadcast with (4,1,3).

score 1 · Answer 3 · answered Jan 23 '20 at 10:47

1

Here is the same thing using numpy's own broadcast function:

def can_be_broadcast(*args):
    try:
        numpy.broadcast(*args)
        return True
    except ValueError:
        return False

answered Jan 23 '20 at 10:47

summentier

456
4
13

Checking if two arrays are broadcastable in python

3 Answers3