我将如何在 C++ 中编写 numpy.tensordot？

Question

我正在尝试用 C++ 复制 numpy.tensordot。 numpy 文档中的示例显示了一个我可以开始工作的嵌套循环，但是如果不是

呢？

c = np.tensordot(a,b, axes=([1,0],[0,1]))

我想做：

c = np.tensordot(a,b, axes=([1,2],[0,1]))

python 中的新嵌套循环会是什么样子？在 C++ 中是否有 easier/faster 方法来执行此操作？现在我在 C++ 中使用与 std::vector 相同的嵌套“for”循环。我看过一些可能有用的库，但我正在尝试仅使用 c++ 标准库。

这是那个 numpy 示例，link 文档：https://numpy.org/doc/stable/reference/generated/numpy.tensordot.html

Examples

A “traditional” example:

>>>
a = np.arange(60.).reshape(3,4,5)
b = np.arange(24.).reshape(4,3,2)
c = np.tensordot(a,b, axes=([1,0],[0,1]))
c.shape
(5, 2)
c
array([[4400., 4730.],
       [4532., 4874.],
       [4664., 5018.],
       [4796., 5162.],
       [4928., 5306.]])
# A slower but equivalent way of computing the same...
d = np.zeros((5,2))
for i in range(5):
  for j in range(2):
    for k in range(3):
      for n in range(4):
        d[i,j] += a[k,n,i] * b[n,k,j]
c == d
array([[ True,  True],
       [ True,  True],
       [ True,  True],
       [ True,  True],
       [ True,  True]])

谢谢

Answer 1

我发现首先重写 np.einsum 很有帮助，因为生成的 for 循环代码在概念上看起来非常相似：

a = np.random.rand(16, 8, 2)
b = np.random.rand(8, 2, 1)


c =  np.tensordot(a, b, axes=([1,2],[0,1]))

# same thing written with einsum
c_ein = np.einsum("ijk,jko->io", a, b)

# same thing done with for loops, 
# notice how we can use the same letters and indexing as einsum
c_manual = np.zeros((16, 1))
for i in range(16):
    for o in range(1):
        # j and k are summed since they don't appear in output
        total = 0
        for j in range(8):
            for k in range(2):
                total += a[i, j, k] * b[j, k, o]
        c_manual[i, o] = total

assert np.allclose(c, c_ein, c_manual)

我将如何在 C++ 中编写 numpy.tensordot？

How would I write numpy.tensordot in c++?

c++

numpy

convolution

tensordot