Python – numpy is added along the first axis

I want to add two arrays of different dimensions by simply performing the same addition along the first axis.

Non-vectorized solution:

x = np.array([[[1,2],[3,4],[5,6]],[[7,8],[9,0],[1,2]],[[3,4],[5,6],[7,8]],[[9,0],[1,2],[3,4]]]) #shape (4,3,2)
y = np.array([[1,2],[3,4],[5,6]]) #shape (3,2)

ans = np.empty(x.shape)
for i in range(x.shape[0]):
    ans[i] = x[i] + y

print(ans) #shape (4,3,2)

How can I make this broadcast appropriate?

Since broadcasting [numpy-doc], you can simply use:

x + y

So here we calculate that the element at index i,j,k is:

x_ijk+y_jk

This gives:

>>> x + y
array([[[ 2,  4],
        [ 6,  8],
        [10, 12]],

[[ 8, 10],
        [12,  4],
        [ 6,  8]],

[[ 4,  6],
        [ 8, 10],
        [12, 14]],

[[10,  2],
        [ 4,  6],
        [ 8, 10]]])
>>> (x + y).shape
(4, 3, 2)

If two arrays are added, for example the first array has three dimensions, the second array has two dimensions, and the last two dimensions of the first left array are equal to the dimensions of the right array, the array on the right side is “broadcast”. This means that it is treated as a three-dimensional array where each subarray is equal to the array on the right.

You can also “introduce” an extra dimension of y anywhere, such as this answer “broadcast” a specific dimension.