Numpy Broadcasting arrays

Question

As I'm trying to understand broadcasting in python, I'm coming across a shape mismatch error. I know this means that the arrays I have don't fit in terms of dimension. My code basically tries to do the following operations on the arrays with the following dimensions:

(256,256,3)*(256,256)+(256,256)

I know the problem is in the multiplication. I was wondering if there is any way to fix this? Can I add an extra dimension to the (256,256) array of the multiplication?

Answer 1

Let's say

A.shape = (256,256,3)
B.shape = (256,256)
C.shape = (256,256)

NumPy broadcasting adds axes on the left by default, so that would result in B and C being broadcasted to

B.shape = (256,256,256)
C.shape = (256,256,256)

and clearly that does not work and is not what you desire, since there is a shape mismatch with A.

So when you want to add an axis on the right , use B[..., np.newaxis] and C[..., np.newaxis] :

A*B[..., np.newaxis] + C[..., np.newaxis]

B[..., np.newaxis] has shape (256,256,1) , which gets broadcasted to (256,256,3) when multiplied with A , and the same goes for C[..., np.newaxis] .

---

B[..., np.newaxis] can also be written as B[..., None] -- since np.newaxis is None . It's a little shorter, but the intent is perhaps not quite as clear.