Numpy - 矩阵乘法返回 ndarray，而不是总和

Question

所有，我有一个应用程序，当两个矩阵相乘时，它需要返回一个 numpy ndarray，而不是一个简单的总和；例如：

import numpy as np
x = np.array([[1, 1, 0], [0, 1, 1]])
y = np.array([[1, 0, 0, 1], [1, 0, 1, 0], [0, 0, 0, 0]])
w = x @ y
>>> array([[2, 0, 1, 1],
           [1, 0, 1, 0]])

但是，要求是返回一个 ndarray（在这种情况下..）：

array([[[1,1,0], [0,0,0], [0,1,0], [1,0,0]],
       [[0,1,0], [0,0,0], [0,1,0], [0,0,0]]])

注意，矩阵乘法操作可能会重复；输出将用作下一个矩阵乘法运算的 ndarrays 的左侧矩阵，这将在第二个矩阵乘法运算后产生一个更高阶的 ndarray，等等。

有什么方法可以实现这一目标吗？我已经研究了重载__add__，并按照此处讨论的那样对__radd__np.ndarray进行了子类化，但主要是遇到了维度不兼容错误。

想法？

更新：

解决@Divakar 的答案 例如，对于链式操作，添加

z = np.array([[1, 1, 0], [0, 0, 0], [1, 0, 0], [0, 1, 0]])
s1 = x[...,None] * y
s2 = s1[...,None] * z

导致不希望的输出。

我怀疑问题始于 s1，在上述情况下，它返回 s1.shape = (2,3,4)。它应该是 (2,4,3)，因为 [2x3][3x4] = [2x4]，但我们并不是在这里真正求和，只是返回一个长度为 3 的数组。

类似地，s2.shape 应该是 (2,3,4,3)，[顺便说一下]它是，但有不想要的输出（这不是'错误'，只是不是我们正在寻找的）。详细说明， s1*z 应该是 [2x4][4x3] = [2x3] 矩阵。矩阵的每个元素本身就是一个 [4x3] 的 ndarray，因为我们在 z 中有 4 行来乘以 s1 中的元素，并且 s1 中的每个元素本身都是 3 个元素长（同样，我们不是算术相加元素，而是返回 ndarrays，其扩展维度是操作的 R 矩阵中的行数。

最终，所需的输出将是：

s2 = array([[[[1, 1, 0],
              [0, 0, 0],
              [0, 1, 0],
              [0, 0, 0]],

              [[1, 1, 0],
               [0, 0, 0],
               [0, 0, 0],
               [1, 0, 0]],

              [[0, 0, 0],
               [0, 0, 0],
               [0, 0, 0],
               [0, 0, 0]]],


             [[[0, 1, 0],
               [0, 0, 0],
               [0, 1, 0],
               [0, 0, 0]],

              [[0, 1, 0],
               [0, 0, 0],
               [0, 0, 0],
               [0, 0, 0]],

              [[0, 0, 0],
               [0, 0, 0],
               [0, 0, 0],
               [0, 0, 0]]]])

Answer 1

将它们扩展到3D并利用broadcasting-

x[:,None] * y.T

或与np.einsum-

np.einsum('ij,jk->ikj',x,y)

路过OP's comment和问题的引述：

... 矩阵乘法操作可以重复；输出将用作下一个矩阵乘法运算的 ndarrays 的左侧矩阵，这将在第二个矩阵乘法运算后产生一个更高阶的 ndarray，等等。

看来，我们需要沿着这些方向做一些事情——

s1 = x[...,None] * y
s2 = s1[...,None] * z # and so on.

虽然，在这种情况下轴的顺序会有所不同，但它似乎是将解决方案扩展到通用数量的传入2D数组的最简单方法。

按照问题中的编辑，似乎您正在从第一个轴开始放置传入数组以进行元素乘法。所以，如果我猜对了，你可以交换轴以获得正确的顺序，就像这样 -

s1c = (x[...,None] * y).swapaxes(1,-1)
s2c = (s1c.swapaxes(1,-1)[...,None] * z).swapaxes(1,-1) # and so on.

如果您只对最终输出感兴趣，请仅在最后阶段交换轴并跳过中间阶段的轴。