来自这个question,我想知道是否可能有更广泛的einsum。让我们假设,我遇到了问题
using PyCall
@pyimport numpy as np
a = rand(10,10,10)
b = rand(10,10)
c = rand(10,10,10)
Q = np.einsum("imk,ml,lkj->ij", a,b,c)
或者类似的东西,如何在没有循环总结的情况下解决这个问题?
最好的问候
答案 0 :(得分:3)
编辑/更新:现在这是一个已注册的软件包,因此您可以Pkg.add("Einsum")
并且您应该很高兴(请参阅下面的示例以开始使用)。
原始答案:我刚刚创建了一些非常初步的代码来执行此操作。它完全遵循马特B.在评论中描述的内容。希望它有所帮助,如果有问题,请告诉我。
https://github.com/ahwillia/Einsum.jl
这是您实现示例的方式:
using Einsum
a = rand(10,10,10)
b = rand(10,10)
c = rand(10,10,10)
Q = zeros(10,10)
@einsum Q[i,j] = a[i,m,k]*b[m,l]*c[l,k,j]
在底层,宏构建了以下一系列嵌套for循环,并在编译之前将它们插入到代码中。 (注意这不是插入的确切代码,它还检查以确保输入的维度一致,使用macroexpand
查看完整代码):
for j = 1:size(Q,2)
for i = 1:size(Q,1)
s = 0
for l = 1:size(b,2)
for k = 1:size(a,3)
for m = 1:size(a,2)
s += a[i,m,k] * b[m,l] * c[l,k,j]
end
end
end
Q[i,j] = s
end
end