Tensorflow：矢量化循环

Question

我必须在Tensorflow中对一个循环进行矢量化。

for i in range(n):
    for j in range(3):
        res[i] += A[i - j + 1, j] * u[i - j + 1]

它是与位置相关的内核的某种卷积。 据说我可以使用tf.einsum完成，但我不知道如何。

我知道的唯一方法是使用tf.while_loop，但我想要更优雅的东西。请帮忙。

Answer 1

由于j上的循环非常小，你可以做一些技巧来“翻转循环”并仅对内循环进行矢量化。这应该足够快。

具体来说，你的循环可以像这样互换：

for j in range(3):
    for i in range(n):
        res[i] += A[i - j + 1, j] * u[i - j + 1]

然后我们可以对内循环进行矢量化（为了清晰起见写出零）

for j in range(3):
    res += A[0 - j + 1:n - j + 1, j] * u[0 - j + 1: n - j + 1]

请注意，对于j = 2，索引将变为负数，这可能表示代码中存在错误。