How is numpy.einsum implemented?

user_11606860 0
Viewed 148

I want to understand how is einsum function in python implemented. I found the source code in numpy/core/src/multiarray/einsum.c.src file but couldn't completely understand it. In particular I want to understand how does it creates the required loops automatically?

For example:

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]

Thank you in advance

ps: This question is not about how to use numpy.einsum

numpy c python
1 Answers

I want to understand how does it creates the required loops automatically?

Well, it does not create the loops the way you think it does. In this case, it creates an iterator operating over multiple arrays and then use it in a generic main loop. In the more general case, there are two main loops: one to iterate over the output array items and one to perform a reduction.

The main function is PyArray_EinsteinSum. In your case, it takes an unoptimized path and end up creating a basic iteration function based on the iterator created previously (ie. iter). This function is get_sum_of_products_function. It basically analyze the einsum operation so to find the best (sum of product) function to call based on a lookup table (like _outstride0_specialized_table). In your specific case, double_sum_of_products_outstride0_two is called. Numpy use a template system so to generate this function automatically at build time (*.c.src files are template files converted to *.c files based on predefined basic comments). In this case, the function is generated from @name@_sum_of_products_outstride0_@noplabel@ and once computed by the C preprocessor it gives something like the following function:

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])));
}

As you can see, there is only one main loop iterating over the previously generated iterator. In your case, stride0 and stride1 are both equal to 8, data0 and data1 are the raw input arrays, dataptr is the raw output array and count is set to 120 initially. Note that the fact both strides are equal to 8 is surprising at first glance since the einsum does not iterate on the two array contiguously. This is because the second array is copied and reorder because Numpy cannot create a uniform view based on the einsum parameters.

Note that the fallback case use for the example code is not particularly optimized and it only produce one value. For example, the much more optimized double_sum_of_products_contig_contig_outstride0_two function can be called from unbuffered_loop_nop2_ndim2 for the following code:

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)

In this case, the double_sum_of_products_contig_contig_outstride0_two performs the reductions for a given output item and unbuffered_loop_nop2_ndim2 iterate over the output array.

If the expression ij, ij -> j is instead used in the above code, then the function double_sum_of_products_contig_two is called which operates the same way than double_sum_of_products_contig_contig_outstride0_two except it reads/writes on the whole output line during the reduction.

user_129395570
Related