Is it possible to express an arbitrary tensor contraction in terms of BLAS routines?

Question

I noticed that libraries like numpy and pytorch are able to perform arbitrary tensor contractions at speeds similar to comparably sized matrix multiplications. This leads me to believe that underneath they somehow express these operations in terms of BLAS routines.

However, I can't think of a general algorithm to do so, especially for arbitrary tensor contractions. This is further complicated if one wanted to avoid tensor transpositions. I'm curious how this is done.

Answer 1

You can use reshape and pointwise multiply to reduce the operation to a matmul in terms of two temporary tensors $c_1,c_2$. Consider following transformations:

1. pointwise mul to turn $a_i * b_i$ into $c_i$
2. outer product to turn $a_i * b_j$ into $c_{ij}$
3. reshape to turn $\sum_{ij} a_{ijk}$ into $\sum_{l}c_{lk}$.

With these operations you can define temporary tensors $a,b$ such that final einsum is of the form $\sum_{ijk}a_{ij}b_{jk}$ (matmul).

Basically, out of the indices which are remain uncontracted, you select some of them to be "left" indices and reshape, that's your $i$ index. Then the rest are your "right" indices, ie $k$. Use reshape to turn the contracted indices into $j$. This determines dimensions of your final matmul.

There's more than one way of doing this, and some are more efficient than others. Existing implementations heuristically pick an order for some common einsums to be a good matmul.

There's a Python package (pip install opt_einsum) which can compute an optimal schedule and print it out.

For instance, you can get the optimal sequence of BLAS operations as follows for einsum $ij,kl,iq,kp$ where all dimensions are size 2:

                                                                                            import torch
import numpy as np
import opt_einsum as oe
einsum_string="ij,kl,iq,kp->ljpq"
views = oe.helpers.build_views(einsum_string, {c: 2 for c in einsum_string})
path, path_info = oe.contract_path(einsum_string, *views, optimize='dp')
print(path_info)

You should see something like this

Complete contraction:  ij,kl,iq,kp->ljpq
         Naive scaling:  6
     Optimized scaling:  4
      Naive FLOP count:  2.560e+2
  Optimized FLOP count:  4.800e+1
   Theoretical speedup:  5.333
  Largest intermediate:  1.600e+1 elements
--------------------------------------------------------------------------------
scaling        BLAS                current                             remaining
--------------------------------------------------------------------------------
   3           GEMM              iq,ij->qj                        kl,kp,qj->ljpq
   3           GEMM              kp,kl->pl                           qj,pl->ljpq
   4   OUTER/EINSUM            pl,qj->ljpq                            ljpq->ljpq

Note that this einsum optimizer finds a schedule which uses two GEMMs followed by an outer product, which is better than what PyTorch/numpy would do by default, which would be to use a single GEMM

See my colab notebook for runnable examples.

For explanations of the theory behind optimal einsum computation, see my write-up on Wolfram Community forums: Tensor networks, Einsum optimizers and the path towards autodiff