Repository: google/TensorNetwork

Issue:
Buggy interplay between ``tf.vectorized_map`` and ``tn.contractors.*``
When a function ``f`` has, say ``tn.contractors.optimal()`` inside, the results of the function is of course right, but the results for ``tf.vectorized_map(f, batched_input)`` might be wrong. 
See a full demo script at  https://gist.github.com/refraction-ray/3a24e1ca09d9c2850d62b40cbf57946c. 
The results are always correct for individual input (jitted or not) and for ``jax.vmap`` interface but become wrong for ``tf.vectorized_map`` interface. I suspect the problem may be related to the interplay of ``tf.vectorized_map`` and the path solver provided by opteinsum, though I cannot see why they are entangled.

The output for the above demo script in tensorflow case is
```bash
$ python vectorized_map_bug.py tensorflow

3.6.8 |Anaconda, Inc.| (default, Dec 30 2018, 01:22:34)
[GCC 7.3.0]
tensornetwork version: 0.4.5
tensorflow version: 2.5.0
tf.Tensor((-0.4161468+0.90929735j), shape=(), dtype=complex64)
tf.Tensor((-0.65364355-0.7568025j), shape=(), dtype=complex64)
jitted
tf.Tensor((-0.4161468+0.90929735j), shape=(), dtype=complex64)
tf.Tensor((-0.65364355-0.7568025j), shape=(), dtype=complex64)
mapfn
tf.Tensor([-0.4161468 +0.90929735j -0.65364355-0.7568025j ], shape=(2,), dtype=complex64)
vectorized map
tf.Tensor([0.99999994+0.j 0.99999994+0.j], shape=(2,), dtype=complex64)
```


Discussion:
Thanks, I'll take a look!
Actually the bug is not necessarily related to ``tn.contractors.*``. I found the same issue persists, even we only use ``contract_between`` in the function:

```python
def u0(inputs):
    """
    compute the following tensor

    node0 -        -       -
      |      gate1   gate2
    node1 -        -       -

    gate1 = gate2 = expm(i*inputs*zz)
    """
    nodes = []
    nodes.append(
        tn.Node(
            np.array(
                [
                    [1.0],
                    [0.0],
                ],
                dtype=npdtype,
            ),
        ),
    )
    bk = nodes[0].backend
    nodes.append(
        tn.Node(
            np.array(
                [
                    [1.0],
                    [0.0],
                ],
                dtype=npdtype,
            ),
        )
    )
    nodes[0][1] ^ nodes[1][1]
    gate1 = tn.Node(
        bk.reshape(
            bk.cos(
                inputs[0],
            )
            * ii
            + bk.sin(inputs[0]) * 1.0j * zz,
            [2, 2, 2, 2],
        )
    )
    gate2 = tn.Node(
        bk.reshape(
            bk.cos(inputs[0]) * ii + bk.sin(inputs[0]) * 1.0j * zz,
            [2, 2, 2, 2],
        )
    )
    nodes.append(gate1)
    nodes.append(gate2)
    gate1[2] ^ nodes[0][0]
    gate1[3] ^ nodes[1][0]
    gate2[2] ^ gate1[0]
    gate2[3] ^ gate1[1]
    s = nodes[0] @ nodes[1]
    s = gate1 @ s
    s = gate2 @ s
    # nn = tn.contractors.optimal(nodes, [gate2[0], gate2[1]])
    return s.tensor[0, 0]


```
I had a look at the script you linked. The following things would be helpful in order to analyse the problem:
1. Could you try to reduce the program to a minimal reproducing program (change to explicitly use tensorflow backend, remove unneeded code, aso)?
2. Explicitly state what the input to the program is (e.g. what types, what dimensions of arrays, aso)
3.  Explicitly state what the program is supposed to compute.
thanks!
@mganahl , thanks for your help. 

1. Actually I have try my best to keep the script a minimal to reproduce the bug (smaller network contraction could have no bug), and the more compact script can be found below with only tf backend support:

```python
import tensornetwork as tn
import numpy as np
import tensorflow as tf

tn.set_default_backend("tensorflow")
npdtype = np.complex64
ii = np.eye(4).astype(npdtype)
zz = np.array([[1, 0, 0, 0], [0, -1, 0, 0], [0, 0, -1, 0], [0, 0, 0, 1]], dtype=npdtype)
inputs = np.array([[1.0], [2.0]], dtype=npdtype)


def u0(inputs):
    """
    compute the following tensor

    node0 -        -       -
      |      gate1   gate2
    node1 -        -       -

    gate1 = gate2 = expm(i*inputs*zz)
    """
    nodes = []
    nodes.append(
        tn.Node(
            np.array(
                [
                    [1.0],
                    [0.0],
                ],
                dtype=npdtype,
            ),
        ),
    )
    nodes.append(
        tn.Node(
            np.array(
                [
                    [1.0],
                    [0.0],
                ],
                dtype=npdtype,
            ),
        )
    )
    nodes[0][1] ^ nodes[1][1]
    gate1 = tn.Node(
        tf.reshape(
            tf.math.cos(inputs[0]) * ii + tf.math.sin(inputs[0]) * 1.0j * zz,
            [2, 2, 2, 2],
        )
    )
    gate2 = tn.Node(
        tf.reshape(
            tf.math.cos(inputs[0]) * ii + tf.math.sin(inputs[0]) * 1.0j * zz,
            [2, 2, 2, 2],
        )
    )
    nodes.append(gate1)
    nodes.append(gate2)
    gate1[2] ^ nodes[0][0]
    gate1[3] ^ nodes[1][0]
    gate2[2] ^ gate1[0]
    gate2[3] ^ gate1[1]
    s = nodes[0] @ nodes[1]
    s = gate1 @ s
    s = gate2 @ s
    return s.tensor[0, 0]


print(u0(inputs[0]))
print(u0(inputs[1]))
print(tf.vectorized_map(u0, inputs))
```

2. the program we use is ``u0``, whose input  is a shape [1] complex64 array,
3. and it computes a tensor contraction with the digram shown in u0's docstring while returns the first component of the final tensor. (It corresponds to a quantum circuit simulation with circuit parameters given by ``inputs`` of ``u0`` and returns the first amplitude of the final wavefunction)


Resolving patch:
diff --git a/tensornetwork/backends/tensorflow/tensorflow_backend.py b/tensornetwork/backends/tensorflow/tensorflow_backend.py
--- a/tensornetwork/backends/tensorflow/tensorflow_backend.py
+++ b/tensornetwork/backends/tensorflow/tensorflow_backend.py
@@ -16,7 +16,6 @@
 from typing import Union
 from tensornetwork.backends import abstract_backend
 from tensornetwork.backends.tensorflow import decompositions
-from tensornetwork.backends.tensorflow import tensordot2
 import functools as fct
 import operator as op
 
@@ -46,7 +45,7 @@ def __init__(self) -> None:
 
   def tensordot(self, a: Tensor, b: Tensor,
                 axes: Union[int, Sequence[Sequence[int]]]) -> Tensor:
-    return tensordot2.tensordot(tf, a, b, axes)
+    return tf.tensordot(a, b, axes)
 
   def reshape(self, tensor: Tensor, shape: Tensor) -> Tensor:
     return tf.reshape(tensor, shape)
@@ -108,7 +107,7 @@ def convert_to_tensor(self, tensor: Tensor) -> Tensor:
     return result
 
   def outer_product(self, tensor1: Tensor, tensor2: Tensor) -> Tensor:
-    return tensordot2.tensordot(tf, tensor1, tensor2, 0)
+    return tf.tensordot(tensor1, tensor2, 0)
 
   #pylint: disable=unused-argument
   def einsum(self,


Test patch: