2D矩阵的“切片”的逐元素乘法以形成3D矩阵

Question

像这样的矩阵乘法

使用numpy在Python中易于实现

import numpy as np
np.array([[1, 2, 3]]) * np.array([[1], [2], [3]])

array([[1, 2, 3],
       [2, 4, 6],
       [3, 6, 9]])

但是在我的情况下，我有2个2D矩阵，我想将它们相乘以形成3D矩阵。实际上，2D矩阵的第一个“切片”是一个数组，我想乘以第二个矩阵的第一个“切片”以形成2D矩阵。对于2D矩阵的所有“切片”，将继续进行此操作。将第一个视为维度[x,z]，将第二个视为维度[y,z]。我想将它们乘以得到[x,y,z]。在numpy中有一种优雅的方法吗？

Answer 1

因为您已经可以将乘法描述为

[x, z] * [y, z] -> [x, y, z]

最直接的解决方案很可能是使用Einsum：

import numpy as np

A = np.arange(12).reshape(4, 3)
# array([[ 0,  1,  2],
#        [ 3,  4,  5],
#        [ 6,  7,  8],
#        [ 9, 10, 11]])
B = np.arange(9).reshape(3, 3)
# array([[0, 1, 2],
#        [3, 4, 5],
#        [6, 7, 8]])

C = np.einsum('xz,yz->xyz', A, B)
# array([[[ 0,  1,  4],
#         [ 0,  4, 10],
#         [ 0,  7, 16]],
# 
#        [[ 0,  4, 10],
#         [ 9, 16, 25],
#         [18, 28, 40]],
# 
#        [[ 0,  7, 16],
#         [18, 28, 40],
#         [36, 49, 64]],
# 
#        [[ 0, 10, 22],
#         [27, 40, 55],
#         [54, 70, 88]]])

另一种选择是简单地使用广播

D = A[:, None, :] * B[None, :, :]
np.allclose(D, C)
# True

Answer 2

在对this StackOverflow question的回复中，我设法弄清楚了。

arr = np.array([[1, 2, 3]])
arr * arr.T

array([[1, 2, 3],
       [2, 4, 6],
       [3, 6, 9]])

mat = np.repeat(arr, 3, axis=0)
mat

array([[1, 2, 3],
       [1, 2, 3],
       [1, 2, 3]])

mat[:,:,None] * np.transpose(mat[:,None,:], axes=(1, 0, 2))

array([[[1, 2, 3],
        [2, 4, 6],
        [3, 6, 9]],

       [[1, 2, 3],
        [2, 4, 6],
        [3, 6, 9]],

       [[1, 2, 3],
        [2, 4, 6],
        [3, 6, 9]]])