numpy.einsum 是如何实现的？

Question

我想了解python中的einsum函数是如何实现的。 我在numpy/core/src/multiarray/einsum.c.src文件中找到了源代码，但无法完全理解。 特别是我想了解它如何自动创建所需的循环？

例如：

import numpy as np
a = np.random.rand(2,3,4,5)
b = np.random.rand(5,3,2,4)

ll = np.einsum('ijkl, ljik ->', a,b) # This should loop over all the 
# four indicies i,j,k,l. How does it create loops for these indices automatically ?

# The assume that under the hood it does the following 
sum1 = 0
for i in range(2):
    for j in range(3):
        for k in range(4):
            for l in range(5):
                sum1 = sum1 + a[i,j,k,l]*b[l,j,i,k]

先感谢您

ps：这个问题不是关于如何使用numpy.einsum

Answer 1

我想了解它如何自动创建所需的循环？

好吧，它不会像您认为的那样创建循环。 在这种情况下，它会创建一个对多个数组进行操作的迭代器，然后在通用主循环中使用它。 在更一般的情况下，有两个主要循环：一个迭代输出数组项，一个执行归约。

主要功能是PyArray_EinsteinSum 。 在您的情况下，它采用未优化的路径并最终基于先前创建的迭代器（即iter ） 创建基本迭代函数。 这个函数是get_sum_of_products_function 。 它基本上分析 einsum 操作，以便根据查找表（如_outstride0_specialized_table ）找到要调用的最佳（乘积之和）函数。 在您的特定情况下， double_sum_of_products_outstride0_two 。 Numpy 使用模板系统，以便在构建时自动生成此函数（*.c.src 文件是根据预定义的基本注释转换为 *.c 文件的模板文件）。 在这种情况下，该函数是从@name@_sum_of_products_outstride0_@noplabel@生成的，并且一旦由 C 预处理器计算，它就会给出类似于以下函数的内容：

static void double_sum_of_products_outstride0_two(int nop, 
                                                    char **dataptr,
                                                    npy_intp const *strides, 
                                                    npy_intp count)
{
    npy_double accum = 0;
    char *data0 = dataptr[0];
    npy_intp stride0 = strides[0];
    char *data1 = dataptr[1];
    npy_intp stride1 = strides[1];

    while (count--)
    {
        accum += (*(npy_double *)data0) * (*(npy_double *)data1);
        data0 += stride0;
        data1 += stride1;
    }

    *((npy_double *)dataptr[2]) = (accum + (*((npy_double *)dataptr[2])));
}

如您所见，只有一个主循环迭代先前生成的迭代器。 在您的情况下， stride0和stride1都等于 8， data0和data1是原始输入数组， dataptr是原始输出数组，并且count最初设置为 120。 请注意，两个步幅都等于 8 的事实乍一看令人惊讶，因为 einsum 不会在两个数组上连续迭代。 这是因为第二个数组被复制并重新排序，因为 Numpy 无法根据 einsum 参数创建统一视图。

请注意，示例代码的后备案例使用并没有特别优化，它只产生一个值。 例如，可以从unbuffered_loop_nop2_ndim2为以下代码调用更优化的double_sum_of_products_contig_contig_outstride0_two函数：

import numpy as np

a = np.random.rand(3, 10)
b = np.random.rand(3, 10)

for i in range(1):
    ll = np.einsum('ij, ij -> i', a, b) 

在这种情况下， double_sum_of_products_contig_contig_outstride0_two对给定的输出项执行归约，并且unbuffered_loop_nop2_ndim2迭代输出数组。

如果在上述代码中改为使用表达式ij, ij -> j ，则调用函数double_sum_of_products_contig_two ，其操作方式与double_sum_of_products_contig_contig_outstride0_two相同，只是它在缩减期间在整个输出行上读取/写入。