Investigating the performance of np.einsum

A reader of my

For Test 2, with optimize=True, the results are drastically different. np.einsum is still slower, but only about 1.5 times slower at worst!

Going Deeper

The release notes of (which itself uses BLAS where appropriate).

Calling tensordot where appropriate would speed up einsum for a matmul-like operation. But we are only the speedup when optimize is True.

So what's going on?

If we read def einsum(*operands, out=None, optimize=False, **kwargs) in almost immediately:

    # If no optimization, run pure einsum

    if optimize is False:

        if specified_out:

            kwargs['out'] = out

        return c_einsum(*operands, **kwargs)

Does c_einsum utilize tensordot? I doubt it. Later on in the code, we see the tensordot call that the .

After a great deal of argument parsing and parameter preparation, the axis iteration order is determined and a special-purpose iterator is prepared. Each yield from the iterator represents a different way to stride over all operands simultaneously.

    /* Allocate the iterator */

    iter = NpyIter_AdvancedNew(nop+1, op, iter_flags, order, casting, op_flags,

                               op_dtypes, ndim_iter, op_axes, NULL, 0);

Assuming certain special-case optimizations don't apply, an appropriate sum-of-products (sop) function is :

        iternext = NpyIter_GetIterNext(iter, NULL);

        if (iternext == NULL) {

            goto fail;

        }

        dataptr = NpyIter_GetDataPtrArray(iter);

        stride = NpyIter_GetInnerStrideArray(iter);

        countptr = NpyIter_GetInnerLoopSizePtr(iter);

        needs_api = NpyIter_IterationNeedsAPI(iter);



        NPY_BEGIN_THREADS_NDITER(iter);

        NPY_EINSUM_DBG_PRINT("Einsum loop\n");

        do {

            sop(nop, dataptr, stride, *countptr);

        } while (!(needs_api && PyErr_Occurred()) && iternext(iter));

That's my understanding of how einsum works, which is admittedly still a little thin - it really deserves more than the hour I've given it. It does generally confirm my suspicions, however, that it acts like a generalized version of the grade-school method of matrix multiplication. Ultimately, it delegates out to a series of "sum of product" operations which rely on "striders" moving through the operands - not too different from what you do with your fingers when you learn matrix multiplication.

Summary

So why is np.einsum faster when you call it with optimize=True? There are two reasons.

The first (and original) reason is it tries to find an optimal contraction path, which results in speedups when more than two operands are involved.

The second (and newer) reason is that when optimize=True, it activates a codepath that calls tensordot where possible, which in turn tries to uses BLAS. And BLAS is about as optimized as it gets.

Therefore, even in the two operand case (where there is only one trivial contraction path possible), optimize makes a big difference, resulting in the substantial speedups we find in our performance tests.