dgasmith / opt_einsum

⚡️Optimizing einsum functions in NumPy, Tensorflow, Dask, and more with contraction order optimization.
https://dgasmith.github.io/opt_einsum/
MIT License
819 stars 67 forks source link

Optimizers for contracting large/complex tensor networks - cotengra #125

Open jcmgray opened 4 years ago

jcmgray commented 4 years ago

As part of this paper I have made a python package - cotengra - that is specifically aimed at contracting large complex tensor networks, mostly in the form of PathOptimizer instances for opt_einsum. I thought

(A) it might be of interest to opt_einsum users, but more practically (B) some of the stuff in it might be useful to incorporate into opt_einsum (though other bits are kinda dependency heavy).

Just to describe some of those 'experimental' bits:

SliceFinder and SlicedContractor see - #95

This is an implementation of contraction slicing (breaking contractions into many lower memory chunks with as little overhead as possible). It's pure python and c.f. #95 probably a good fit for opt_einsum. Anyhow if anyone wants to try it and/or suggest improvements they now can!

A ContractionTree object

This 'tree' is really the core object that says how much space and time a contraction will take. It has a few nice advantages compared to handling an explicit path + pathinfo:

  1. Less redundancy - e.g. the trees for path=[(0, 1), (0, 1), (0, 1)] and [(2, 3), (0, 1), (0, 1)] are the same. As well as then being able to check tree1 == tree2 this allows the terms to be rearranged so that as many constant tensors come last, memory is best etc.
  2. You can combine lots of methods to build a single path, i.e. some initial simplifications (#114), some optimal subtrees, some heuristic bits etc.
  3. With a few tweaks it might be a fast way to generate (or replicate the same information as) the PathInfo - i.e. figuring out remaining indices etc.
  4. Its very easy to generate visualisations from:

test3d

On the other hand these are kinda low-level and not really killer features if you just want to perform contractions.

HyperOptimizer stuff

This is an extension of the RandomOptimizer, but where you use Bayesian optimization / Gaussian processes to tune any algorithmic parameters and select between multiple possible path algorithms - so like guided sampling.

For example, the default mode selects between a tunable greedy (bottom-up) and a tunable hypergraph partitioning (top-down) method, to achieve very good performance across lots of different types of tensor expression.

It needs quite a lot of dependencies and is fairly heavy to run (basically good only if you want v HQ contraction paths for large TNs), so maybe best as a separate package.

dgasmith commented 2 years ago

@jcmgray these seem like great features, but often require networkx or other dependancies. Due to the install base, it doesn't seem like we can reasonably require anything beyond NumPy. Are there features that do not have extra depends that you think are a good candidate?