np.einsum 中省略号的一般使用

Question

我正在尝试使用

np.einsum

实现张量网络类型的计算。 计算的目标是让局部 2x2 矩阵顺序作用于大小为 (2,2,2,...,2,2,2) 的张量，其中维度数可以是任意的（只要内存允许） ）。 我对大小为 (2,2,2) 的张量有一个简单的实现

X = np.array([[1,2],[3,4]])
Y = np.ones(2**3).reshape((2,2,2))
Aijk = np.einsum('zi,zjk', X, Y)
Bijk = np.einsum('zi,jzk', X, Aijk)
Cijk = np.einsum('zi,jkz', X, Bijk)

可以看出，对于后续的每次计算，要求和的索引应该向右移动。我可以生成适当的显式字符串，但我想知道是否可以使用省略号来实现它，但我不确定如果索引不是最左边或最右边的索引，如何指定索引位置。

X = np.array([[1,2],[3,4]])
Y = np.ones(2**n).reshape(shape) # where shape is the appropriate n-tuple
A = np.einsum('a...,a...', X, Y)
B = np.einsum('b...,...b...', X, A) # how do I specify that the summed index of A is the second one?
...
N = np.einsum('l...,...l', X, M)

我希望我想要实现的目标很清楚。

Answer 1

让我们采取一些小步骤在您的示例中添加省略号。

In [3]: X = np.array([[1,2],[3,4]])
   ...: Y = np.ones(2**3).reshape((2,2,2))
   ...: Aijk = np.einsum('zi,zjk', X, Y)

In [4]: Aijk
Out[4]: 
array([[[4., 4.],
        [4., 4.]],

       [[6., 6.],
        [6., 6.]]])

使 RHS 明确：

In [5]: np.einsum('zi,zjk->ijk', X, Y)
Out[5]: 
array([[[4., 4.],
        [4., 4.]],

       [[6., 6.],
        [6., 6.]]])

将重复的

jk

替换为省略号：

In [6]: np.einsum('zi,z...->i...', X, Y)
Out[6]: 
array([[[4., 4.],
        [4., 4.]],

       [[6., 6.],
        [6., 6.]]])

并尝试使用其他

Y

形状（始终保持第一维度，

2

来匹配

z

：

In [7]: np.einsum('zi,z...->i...', X, np.ones((2,2,2,2)))
Out[7]: 
array([[[[4., 4.],
         [4., 4.]],

        [[4., 4.],
         [4., 4.]]],


       [[[6., 6.],
         [6., 6.]],

        [[6., 6.],
         [6., 6.]]]])

In [12]: np.einsum('zi,z...->i...', X, np.ones((2,3)))
Out[12]: 
array([[4., 4., 4.],
       [6., 6., 6.]])

In [13]: np.einsum('zi,z...->i...', X, np.ones((2,1,2,3)))
Out[13]: 
array([[[[4., 4., 4.],
         [4., 4., 4.]]],


       [[[6., 6., 6.],
         [6., 6., 6.]]]])

我想我可以使用 'zi,z...->...i', 'zi,jz...->ij...', 'zi,jz...->i...j '。 （但我还没有测试过它们）。但我很确定我不能使用“...i...”。并且一般不能使用省略号来指定求和“z”维度的相对位置。