Python numpy.einsum 是如何实现的

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

好吧，它不会像你想象的那样创建循环.在这种情况下，它创建了一个在多个数组上运行的迭代器，然后在通用主循环中使用它.在更一般的情况下，有两个主要循环:一个循环迭代输出数组项，另一个循环执行缩减.

主要功能是PyArray_EinsteinSum.在您的例子中，它采用了一条未优化的路径，并基于之前创建的迭代器(即iter)以creating a basic iteration function结尾.此函数为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迭代输出array.

如果在上述代码中改用表达式ij, ij -> j，则调用函数double_sum_of_products_contig_two，该函数的操作方式与double_sum_of_products_contig_contig_outstride0_two相同，只是在归约期间读取/写入整个输出行.