TypeError when trying to run the notebook

[bgavran]

When trying to run the Jupyter Notebook, the small scale example, I get the following error:

Found 2 einsums.
The running time of the whole function is 0.000613 s
Einsums constitue 131.8 % of the running time (0.000807 s).
---------------------------------------------------------------------------
TypeError                                 Traceback (most recent call last)
<ipython-input-8-8750a51928fb> in <module>()
      7 c = tf.random_normal((12, 14))
      8 # res = func(a, b, c)
----> 9 orders, optimized_func = tf_einsum_opt.optimizer(func, sess, a, b, c)

/home/bgavran3/Dropbox/ucenje/Programiranje/tf_einsum_opt/tf_einsum_opt.py in optimizer(f, sess, *args)
    115   improved_orders = {}
    116   for idx in range(len(slowest_to_fastest)):
--> 117     caller_str = cache.keys()[slowest_to_fastest[idx]]
    118     vanilla_einsum_timings = cache[caller_str]['timings']
    119     rel_timing = np.sum(vanilla_einsum_timings) / vanilla_whole_runtime

TypeError: 'dict_keys' object does not support indexing

I'm using Python 3.6 and Tensorflow 1.2.0rc1, not sure if any other info is needed as I've just cloned the repo and ran the notebook.

[Bihaqo]

Hi, thanks for the issue!

Should work now. I have a nasty habit of using Python 2 in new projects and I haven't even tested this one in Python 3, shame on me.
Note that for some reason in Python 3 (at least in TF 1.1) the small scale example does bring any improvements in speed, probably the speed difference between the default order of execution and the optimal order is too small to catch. Try the large scale one, but the optimization will take ~10 minutes.

BTW, after some thinking, I started to believe that this optimization should be done analytically (so much faster and a little bit less reliable) during the graph compilation time. I will try it out and try to pull request into TF.

[bgavran]

Yes, it works now!

I agree, although I'm not sure I understand the problems behind analytical optimizations of computational graph.
I asked a similar question on reddit a while ago and apparently, it is not trivial.

[Bonnevie]

I believe that einsum contractions are identical to contractions in so-called tensor networks. This article
https://arxiv.org/abs/1710.01437
shows (as far as I understand it) that the optimal contraction order is given by mapping the tensor network to a graphical model and then running message-passing/the junction tree algorithm from the graphical models literature. Provides an avenue of attack and is perhaps a bit more palatable if you are from an ML background.