Numpy Broadcast Addition Along Arbitrary Axes

I would like to add two arrays with different dimensions by simply performing an identical addition along one or more axes. A non-vectorized solution: x = np.array([[[1,2],[3,4],[5

Solution 1:

Look at the shapes of the two arrays:

>>> x.shape
(4, 3, 2)
>>> y.shape
(4, 2)

You see the addition will need to be broadcasted along the 0th and last axis here. A simple option would be

>>> x + y[:, None, :] 
array([[[ 2,  4],
        [ 4,  6],
        [ 6,  8]],

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

       [[ 8, 10],
        [10, 12],
        [12, 14]],

       [[16,  8],
        [ 8, 10],
        [10, 12]]])

Where,

>>> y[:, None, :].shape
(4, 1, 2)

Which effectively just changes the strides of y so the addition can be broadcasted.

---

Better still, use np.expand_dims as suggested by hpaulj in the comments, this'll add an extra penultimate dimension, so you could just do

>>> x + np.expand_dims(y, 1)
array([[[ 2,  4],
        [ 4,  6],
        [ 6,  8]],

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

       [[ 8, 10],
        [10, 12],
        [12, 14]],

       [[16,  8],
        [ 8, 10],
        [10, 12]]])