numpy broadcasting with 3d arrays

Question

Is it possible to apply numpy broadcasting (with 1D arrays),

x=np.arange(3)[:,np.newaxis]
y=np.arange(3)
x+y=
array([[0, 1, 2],
       [1, 2, 3],
       [2, 3, 4]])

to 3d matricies similar to the one below, such that each element in a[i] is treated as a 1D vector like in the example above?

a=np.zeros((2,2,2))
a[0]=1
b=a
result=a+b

resulting in

result[0,0]=array([[2, 2],
                   [2, 2]])

result[0,1]=array([[1, 1],
                   [1, 1]])

result[1,0]=array([[1, 1],
                   [1, 1]])

result[1,1]=array([[0, 0],
                   [0, 0]])

Answer 1

You can do this in the same way as if they are 1d array, ie, insert a new axis between axis 0 and axis 1 in either a or b :

a + b[:,None]    # or a[:,None] + b

---

(a + b[:,None])[0,0]
#array([[ 2.,  2.],
#       [ 2.,  2.]])

(a + b[:,None])[0,1]
#array([[ 1.,  1.],
#       [ 1.,  1.]])

(a + b[:,None])[1,0]
#array([[ 1.,  1.],
#       [ 1.,  1.]])

(a + b[:,None])[1,1]
#array([[ 0.,  0.],
#       [ 0.,  0.]])

Answer 2

Since a and b are of same shape, say (2,2,2) , a+b will indeed work. The way broadcasting works is that it matches the dimensions of the operands in reverse order, starting from the last dimension going up (eg considering columns before rows in a two-dimensional case). If the dimensions match then the next dimension is considered.

In case the dimensions don't match AND if one of the dimensions is 1 then that operand's dimension is repeated to match the other operand (eg if a.shape = (2,1,2) and b.shape = (2,2,2) then the values at the 1st dimension of a are repeated to make the shape (2,2,2) )