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

Answer 1

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

Answer 2

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).

Answer 3

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