Умножение массива 3D numpy на массив 2D numpy

Question

Мне нужно умножить массив 3D numpy на массив 2D numpy.

Допустим, трехмерный массив A имеет форму (3, 100, 500), а двумерный массив B имеет форму (3, 100). Мне нужно поэлементное умножение для каждой из этих 500 осей в трехмерном массиве на двумерный массив, а затем мне нужно сложить по первой оси результирующего массива, получая массив размером (100, 500).

Я могу получить пару циклов for, но наверняка должна быть функция numpy, которая будет достигать этого в 1 строке? Я посмотрел на np.tensordot, np.dot, np.matmul, np.prod и np.sum, но ни одна из этих функций не сделает именно это.

Answer 1

Вы можете попробовать следующее, которое должно работать,

np.sum(A.T*B.T,axis=-1).T

Это даст вам форму (100,500)

Answer 2

Мы можем использовать numpy вещание:

import numpy as np

a = np.random.rand(3,100,500)
b = np.random.rand(3,100)

# add new axis to b to use numpy broadcasting
b = b[:,:,np.newaxis]
#b.shape = (3,100,1)

# elementwise multiplication
m = a*b
# m.shape = (3,100,500)

# sum over 1st axis
s = np.sum(m, axis=0)

#s.shape = (100,500)

Answer 3

Вы можете осуществлять вещание, добавляя новую единичную ось в двумерный массив:

np.sum(A * B[..., None], axis=0)

None в индексе вводит единичный размер в этой позиции, который можно использовать для выравнивания осей для вещания. ... является сокращением для : столько раз, сколько существует измерений: в этом случае оно эквивалентно :, :, поскольку B является 2D.

Альтернативным способом записи будет

(A * B.reshape(*B.shape, 1)).sum(axis=0)

Answer 4

Вы можете легко express выполнить эти операции, используя np.einsum, в данном случае:

np.einsum("ijk,ij->jk", A, B)