朱莉娅的numpy.einsum?(2)

2022-01-23 00:00:00 numpy julia numpy-einsum

问题描述

来自这个问题，我想知道是否有可能更广义的einsum.让我们假设，我有问题

Coming from this question, I wonder if a more generalized einsum was possible. Let us assume, I had the problem

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)

或者类似的东西，我如何在不循环求和的情况下解决这个问题?

Or something similar, how were I to solve this problem without looping through the sums?

致以最诚挚的问候

解决方案

编辑/更新: 现在这是一个注册包，所以你可以 Pkg.add("Einsum")，您应该一切顺利(请参阅下面的示例以开始使用).

Edit/Update: This is now a registered package, so you can Pkg.add("Einsum") and you should be good to go (see the example below to get started).

原答案:我刚刚创建了一些非常初步的代码来执行此操作.它完全符合马特 B. 在他的评论中描述的内容.希望对你有帮助，如果有问题请告诉我.

Original Answer: I just created some very preliminary code to do this. It follows exactly what Matt B. described in his comment. Hope it helps, let me know if there are problems with it.

https://github.com/ahwillia/Einsum.jl

这就是你将如何实现你的例子:

This is how you would implement your example:

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 查看完整代码):

Under the hood the macro builds the following series of nested for loops and inserts them into your code before compile time. (Note this is not the exact code inserted, it also checks to make sure the dimensions of the inputs agree, using macroexpand to see the full code):

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

相关文章