numpy.einsum 具有不同维度的椭圆答案

【问题标题】：numpy.einsum with ellipses of different dimensionalitynumpy.einsum 具有不同维度的椭圆
【发布时间】：2022-11-09 21:22:38
【问题描述】：

我经常发现我想在两个数组的最后几个维度之间进行操作，其中第一个维度不一定匹配。作为一个例子，我想做类似的事情：

a = np.random.randn(10, 10, 3, 3)
b = np.random.randn(5, 3)
c = np.einsum('...ij, ,,,j -> ...,,,i', a, b)

并且结果应该满足c.shape = (10, 10, 5, 3) 和c[i, j, k] = a[i, j] @ b[k]。有没有办法用现有的接口来实现这一点？

【问题讨论】：

这些逗号是怎么回事？
np.einsum('...ij,kj->...ki', a, b)，多了一个可区分的轴。

标签： numpy array-broadcasting numpy-einsum

【解决方案1】：

In [82]: c = np.einsum('...ij,...j->...i', a, b)
---------------------------------------------------------------------------
ValueError                                Traceback (most recent call last)
Input In [82], in <cell line: 1>()
----> 1 c = np.einsum('...ij,...j->...i', a, b)

File <__array_function__ internals>:5, in einsum(*args, **kwargs)

File ~naconda3libsite-packages
umpycoreeinsumfunc.py:1359, in einsum(out, optimize, *operands, **kwargs)
   1357     if specified_out:
   1358         kwargs['out'] = out
-> 1359     return c_einsum(*operands, **kwargs)
   1361 # Check the kwargs to avoid a more cryptic error later, without having to
   1362 # repeat default values here
   1363 valid_einsum_kwargs = ['dtype', 'order', 'casting']

ValueError: operands could not be broadcast together with remapped shapes 
[original->remapped]: (10,10,3,3)->(10,10,3,3) (5,3)->(5,newaxis,3)

所以它试图使用broadcasting 来匹配尺寸。

让我们制作a (10,10,1,3,3) 形状。这样 (10,10,1) 部分将与 b 的 (5,) 一起广播：

In [83]: c = np.einsum('...ij,...j->...i', a[:,:,None], b)
In [84]: c.shape
Out[84]: (10, 10, 5, 3)

【讨论】：

【解决方案2】：

最终解决了以下帮助函数的问题

def batch_matvec(A, b):
    product = np.einsum('...ij, ...kj->...ki', A, b.reshape(-1, b.shape[-1]))
    return product.reshape((*A.shape[:-2], *b.shape[:-1], -1))

【讨论】：