`numpy.einsum` 中的 `out` 参数无法按预期工作

Posted 2023-03-13

【中文标题】`numpy.einsum` 中的 `out` 参数无法按预期工作

【英文标题】：The `out` arguments in `numpy.einsum` can not work as expected

【发布时间】：2018-05-09 14:13:39

【问题描述】：

我有两段代码。第一个是：

A = np.arange(3*4*3).reshape(3, 4, 3)
P = np.arange(1, 4)
A[:, 1:, :] = np.einsum('j, ijk->ijk', P, A[:, 1:, :])

结果A 是：

array([[[  0,   1,   2],
        [  6,   8,  10],
        [ 18,  21,  24],
        [ 36,  40,  44]],

       [[ 12,  13,  14],
        [ 30,  32,  34],
        [ 54,  57,  60],
        [ 84,  88,  92]],

       [[ 24,  25,  26],
        [ 54,  56,  58],
        [ 90,  93,  96],
        [132, 136, 140]]])

第二个是：

A = np.arange(3*4*3).reshape(3, 4, 3)
P = np.arange(1, 4)
np.einsum('j, ijk->ijk', P, A[:, 1:, :], out=A[:,1:,:])

结果A 是：

array([[[ 0,  1,  2],
        [ 0,  0,  0],
        [ 0,  0,  0],
        [ 0,  0,  0]],

       [[12, 13, 14],
        [ 0,  0,  0],
        [ 0,  0,  0],
        [ 0,  0,  0]],

       [[24, 25, 26],
        [ 0,  0,  0],
        [ 0,  0,  0],
        [ 0,  0,  0]]])

所以结果是不同的。这里我想使用out来节省内存。这是numpy.einsum 中的错误吗？还是我错过了什么？

顺便说一下，我的numpy 版本是1.13.3。

【参考方案1】：

我之前没有使用过这个新的out 参数，但过去使用过einsum，并且大致了解它的工作原理（或至少曾经使用过）。

在我看来，它在迭代开始之前将 out 数组初始化为零。这将解释 A[:,1:,:] 块中的所有 0。相反，如果我初始单独的 out 数组，则插入所需的值

In [471]: B = np.ones((3,4,3),int)
In [472]: np.einsum('j, ijk->ijk', P, A[:, 1:, :], out=B[:,1:,:])
Out[472]: 
array([[[  3,   4,   5],
        [ 12,  14,  16],
        [ 27,  30,  33]],

       [[ 15,  16,  17],
        [ 36,  38,  40],
        [ 63,  66,  69]],

       [[ 27,  28,  29],
        [ 60,  62,  64],
        [ 99, 102, 105]]])
In [473]: B
Out[473]: 
array([[[  1,   1,   1],
        [  3,   4,   5],
        [ 12,  14,  16],
        [ 27,  30,  33]],

       [[  1,   1,   1],
        [ 15,  16,  17],
        [ 36,  38,  40],
        [ 63,  66,  69]],

       [[  1,   1,   1],
        [ 27,  28,  29],
        [ 60,  62,  64],
        [ 99, 102, 105]]])

einsum 的 Python 部分并没有告诉我太多，除了它决定如何将 out 数组传递给 c 部分（作为 tmp_operands 的列表之一）：

c_einsum(einsum_str, *tmp_operands, **einsum_kwargs)

我知道它设置了 c-api 等效于 np.nditer，使用 str 定义轴和迭代。

它在迭代教程中迭代了类似于本节的内容：

https://docs.scipy.org/doc/numpy-1.13.0/reference/arrays.nditer.html#reduction-iteration

特别注意it.reset() 步骤。这会在迭代之前将 out 缓冲区设置为 0。然后它遍历输入数组和输出数组的元素，将计算值写入输出元素。由于它正在做一个产品总和（例如out[:] += ...），它必须从头开始。

我对实际发生的情况进行了一些猜测，但对我来说，它应该从零开始输出缓冲区似乎是合乎逻辑的。如果该数组与输入之一相同，则最终会弄乱计算。

所以我认为这种方法不会有效并节省您的记忆。它需要一个干净的缓冲区来累积结果。完成后，或者您可以将值写回A。但鉴于 dot 类似产品的性质，您不能将相同的数组用于输入和输出。

In [476]: A[:,1:,:] = np.einsum('j, ijk->ijk', P, A[:, 1:, :])
In [477]: A
Out[477]: 
array([[[  0,   1,   2],
        [  3,   4,   5],
        [ 12,  14,  16],
        [ 27,  30,  33]],
        ....)

【讨论】：

hpaulj，@ely 你们中的任何一个都知道这是否（或应该）以任何方式与这个新的“Inplace 操作检查输入是否重叠输出并创建临时以避免问题。 " numpy1.13 中的功能？

@PaulPanzer，我从发行说明中怀疑+= 类似操作和相关ufunc 的处理发生了变化。 np.add(fib[:-2], fib[1:-1], out=fib[2:])、***.com/questions/47427603/recurrence-with-numpy，给出不同的结果 dtype 和版本 (12 v 13)。在13 中，看起来整个数组都被缓冲了，并且没有递归。但即使有out 参数，einsum 也不是ufunc。

【参考方案2】：

在einsum、there is a section 的 C 源代码中，它将采用 out 指定的数组并进行一些零设置。

但以Python source code 为例，有一些执行路径在调用c_einsum 之前调用tensordot 函数。

这意味着某些操作可能是使用tensordot 预先计算的（因此在某些收缩过程中修改您的数组A），在任何子数组都被设置为零之前einsum 的 C 代码中的零设置器。

另一种说法是：在每次执行下一个收缩操作时，NumPy 都有很多可用的选择。要直接使用tensordot 而无需进入C 级einsum 代码吗？还是准备参数并传递到 C 级别（这将涉及用全零覆盖输出数组的某些子视图）？还是重新排序操作并重复检查？

根据它为这些优化选择的顺序，您最终可能会得到意想不到的全零子数组。

最好的办法是不要尝试变得如此聪明并使用相同的数组作为输出。你说这是因为你想节省内存。是的，在某些特殊情况下，einsum 操作可能就地可行。但它目前不检测是否是这种情况并尝试避免置零。

在很多情况下，在整个操作的中间覆盖输入数组之一会导致很多问题，就像尝试附加到您直接循环的列表等一样。