Checking if two arrays are broadcastable in python

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

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

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