добавление широковещательной трансляции по произвольным осям

Question

Я хотел бы добавить два массива с разными размерами, просто выполнив одинаковое сложение вдоль одной или нескольких осей.

Невекторное решение:

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],[7,8]]) #shape (4,2)

ans = np.copy(x)
for i in range(x.shape[1]):
    ans[:,i] += y

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

Как я могу транслировать это вдоль (например) второй оси? В целом, как я могу сделать эту трансляцию по произвольным осям?

Answer 1

Посмотрите на формы двух массивов:

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

Вы видите, что дополнение нужно будет транслировать вдоль 0-й и последней оси здесь. Простой вариант будет

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

Где,

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

Что фактически просто изменяет шаги y, так что добавление может транслироваться.

---

Еще лучше, используйте np.expand_dims, как предложено hpaulj в комментариях, это добавит дополнительное предпоследнее измерение, так что вы можете просто сделать

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