Home

Awesome

OMEinsum-Benchmarks

This repository is used to compare the performance of einsum in python (numpy, torch), with implementations in julia.

In /pythonscripts, the benchmarking-scripts for numpy.einsum and torch.einsum can be found. The tested operations and the sizes of the arrays involved are:

labeldescriptioneinsum-equationsmallmediumlarge
matmulmatrix multiplicationij,jk -> ik10×10100×1001000×1000
dotdot productijl,ijl->10×10×1020×20×2050×50×50
ptracepartial traceiij->j10×10×1020×20×2050×50×50
hadamardthe hadamard productijk,ijk->ijk10×10×1030×30×3050×50×50
tcontracttensor contractionikl,kjl->ij2×2×210×10×1030×30×30
indexsumsumming over a (single) indexijk->ik10×10×1030×30×30100×100×100
tracematrix-traceii->10×10100×1001000×1000
diagtaking the diagonal w.r.t to two or more indicesjii->ji10×10×1020×20×2050×50×50
starandcontractmixing tensor and star-contractionij,il,il->j10×1030×30100×100
starstarcontracting over a single indexij,ik,il->jkl10×1030×30100×100
outertaking the outer product between two tensorsij,kl->ijkl10×1050×50100×100
batchmulbatchmultiplying two batches of matricesijl,jkl->ikl10×10×3100×100×31000×1000×3
permpermuting a tensors indicesijkl->ljki2×2×2×210×10×10×1030×30×30×30

The benchmarks have been evaluated over inputs of different size, labelled (ad-hoc) as tiny, small, medium, large and huge. The current benchmark has been run with torch, numpy for python and OMEinsum.jl and Einsum.jl for julia.

Overview

The plots are split by the memory footprint of the operations, i.e. by size size of the tensors involved which have been classified as tiny, small, medium, large or huge rather arbitrarily. The results below are for tensors of type float64 but more results can be found in the plots/ folder. Generally, changing to float32 or complex types does not change the picture too much, except that torch does not support complex numbers.

small tensors

medium tensors

large tensors