Repository: google/TensorNetwork

Issue:
ncon_interface tests fail 



Discussion:

Resolving patch:
diff --git a/conftest.py b/conftest.py
--- a/conftest.py
+++ b/conftest.py
@@ -16,9 +16,33 @@
 from __future__ import division
 from __future__ import print_function
 import pytest
+import jax
+import tensornetwork
+import tensorflow as tf
 
 
 @pytest.fixture(name="backend", params=["numpy", "tensorflow",
                                         "jax", "pytorch"])
 def backend_fixure(request):
   return request.param
+
+
+@pytest.fixture(autouse=True)
+def reset_default_backend():
+  tensornetwork.set_default_backend("numpy")
+  yield
+  tensornetwork.set_default_backend("numpy")
+
+
+@pytest.fixture(autouse=True)
+def enable_jax_64():
+  jax.config.update("jax_enable_x64", True)
+  yield
+  jax.config.update("jax_enable_x64", True)
+
+
+@pytest.fixture(autouse=True)
+def tf_enable_v2_behaviour():
+  tf.compat.v1.enable_v2_behavior()
+  yield
+  tf.compat.v1.enable_v2_behavior()
diff --git a/tensornetwork/__init__.py b/tensornetwork/__init__.py
--- a/tensornetwork/__init__.py
+++ b/tensornetwork/__init__.py
@@ -12,4 +12,4 @@
 
 def set_default_backend(backend: Text, dtype: Optional[Type] = None) -> None:
   config.default_backend = backend
-  config.default_dype = dtype
+  config.default_dtype = dtype


Test patch:
diff --git a/experiments/MERA/binary_mera_test.py b/experiments/MERA/binary_mera_test.py
--- a/experiments/MERA/binary_mera_test.py
+++ b/experiments/MERA/binary_mera_test.py
@@ -18,12 +18,16 @@
 import numpy as np
 import tensornetwork as tn
 import experiments.MERA.binary_mera_lib as bml
-import experiments.MERA.binary_mera as bm
 import experiments.MERA.misc_mera as misc_mera
 import pytest
-import copy
-tn.set_default_backend("tensorflow")
-tf.enable_v2_behavior()
+
+
+@pytest.fixture(autouse=True)
+def backend_tensorflow():
+    tn.set_default_backend("tensorflow")
+    yield
+    tn.set_default_backend("numpy")
+
 
 @pytest.mark.parametrize("chi", [4, 6])
 @pytest.mark.parametrize("dtype", [tf.float64, tf.complex128])
diff --git a/experiments/MPS/dmrg_test.py b/experiments/MPS/dmrg_test.py
--- a/experiments/MPS/dmrg_test.py
+++ b/experiments/MPS/dmrg_test.py
@@ -15,16 +15,12 @@
 unittests
 """
 import tensorflow as tf
-import tensornetwork as tn
-import experiments.MPS.misc_mps as misc_mps
 import experiments.MPS.matrixproductstates as MPS
 import experiments.MPS.DMRG as DMRG
 import experiments.MPS.matrixproductoperators as MPO
 import pytest
 import numpy as np
 
-tf.enable_v2_behavior()
-
 
 @pytest.mark.parametrize("dtype", [tf.float64, tf.complex128])
 def test_iDMRG(dtype):
diff --git a/experiments/MPS/mps_test.py b/experiments/MPS/mps_test.py
--- a/experiments/MPS/mps_test.py
+++ b/experiments/MPS/mps_test.py
@@ -15,16 +15,11 @@
 unittests
 """
 import tensorflow as tf
-import tensornetwork as tn
 import experiments.MPS.misc_mps as misc_mps
 import experiments.MPS.matrixproductstates as MPS
-import experiments.MPS.DMRG as DMRG
-import experiments.MPS.matrixproductoperators as MPO
 import pytest
 import numpy as np
 
-tf.enable_v2_behavior()
-
 
 @pytest.mark.parametrize("direction", ['l', 'r'])
 def test_TMeigs_float(direction):
diff --git a/experiments/MPS_classifier/MPSMNIST_test.py b/experiments/MPS_classifier/MPSMNIST_test.py
--- a/experiments/MPS_classifier/MPSMNIST_test.py
+++ b/experiments/MPS_classifier/MPSMNIST_test.py
@@ -1,12 +1,8 @@
-import sys
-from sys import stdout
 import numpy as np
 import experiments.MPS.misc_mps as misc_mps
-import itertools
 import experiments.MPS_classifier.MPSMNIST as mm
 import pytest
 import tensorflow as tf
-tf.enable_v2_behavior()
 
 
 def test_predict():
diff --git a/experiments/MPS_classifier/batchtensornetwork_test.py b/experiments/MPS_classifier/batchtensornetwork_test.py
--- a/experiments/MPS_classifier/batchtensornetwork_test.py
+++ b/experiments/MPS_classifier/batchtensornetwork_test.py
@@ -1,59 +1,56 @@
-# Copyright 2019 The TensorNetwork Authors
-#
-# Licensed under the Apache License, Version 2.0 (the "License");
-# you may not use this file except in compliance with the License.
-# You may obtain a copy of the License at
-#
-#      http://www.apache.org/licenses/LICENSE-2.0
-#
-# Unless required by applicable law or agreed to in writing, software
-# distributed under the License is distributed on an "AS IS" BASIS,
-# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
-# See the License for the specific language governing permissions and
-# limitations under the License.
-
-from __future__ import absolute_import
-from __future__ import division
-from __future__ import print_function
-
-import pytest
-import numpy as np
-import tensorflow as tf
-from experiments.MPS_classifier import batchtensornetwork
-
-tf.enable_v2_behavior()
-
-
-@pytest.fixture(
-  name="backend", params=["numpy", "tensorflow", "jax"])
-def backend_fixure(request):
-    return request.param
-
-def test_batched_contract_between_with_matrices(backend):
-  net = batchtensornetwork.BatchTensorNetwork(backend=backend)
-  a = net.add_node(np.ones([10, 2, 2]))
-  b = net.add_node(np.ones([10, 2, 2]))
-  net.connect(a[2], b[1])
-  c = net.batched_contract_between(a, b, a[0], b[0])
-  np.testing.assert_allclose(c.tensor, 2 * np.ones([10, 2, 2]))
-  
-def test_batched_contract_between_multiple_shared_edges(backend):
-  net = batchtensornetwork.BatchTensorNetwork(backend=backend)
-  a = net.add_node(np.ones([5, 3, 4, 9]))
-  b = net.add_node(np.ones([9, 4, 5]))
-  net.connect(a[2], b[1])
-  net.connect(a[3], b[0])
-  c = net.batched_contract_between(a, b, a[0], b[2])
-  np.testing.assert_allclose(c.tensor, 36 * np.ones([5, 3]))
-  
-def test_pairwise_reduction():
-  net = batchtensornetwork.BatchTensorNetwork(backend="tensorflow")
-  a = net.add_node(np.ones([10, 2, 2]))
-  b = batchtensornetwork.pairwise_reduction(net, a, a[0])
-  np.testing.assert_allclose(b.tensor, 512 * np.ones([2, 2]))
-  
-def test_pairwise_reduction_multiple_edges():
-  net = batchtensornetwork.BatchTensorNetwork(backend="tensorflow")
-  a = net.add_node(np.ones([3, 5, 6, 2, 2]))
-  b = batchtensornetwork.pairwise_reduction(net, a, a[1])
-  np.testing.assert_allclose(b.tensor, 16 * np.ones([3, 6, 2, 2]))
+# Copyright 2019 The TensorNetwork Authors
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#      http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+from __future__ import absolute_import
+from __future__ import division
+from __future__ import print_function
+
+import pytest
+import numpy as np
+from experiments.MPS_classifier import batchtensornetwork
+
+
+@pytest.fixture(
+  name="backend", params=["numpy", "tensorflow", "jax"])
+def backend_fixure(request):
+    return request.param
+
+def test_batched_contract_between_with_matrices(backend):
+  net = batchtensornetwork.BatchTensorNetwork(backend=backend)
+  a = net.add_node(np.ones([10, 2, 2]))
+  b = net.add_node(np.ones([10, 2, 2]))
+  net.connect(a[2], b[1])
+  c = net.batched_contract_between(a, b, a[0], b[0])
+  np.testing.assert_allclose(c.tensor, 2 * np.ones([10, 2, 2]))
+  
+def test_batched_contract_between_multiple_shared_edges(backend):
+  net = batchtensornetwork.BatchTensorNetwork(backend=backend)
+  a = net.add_node(np.ones([5, 3, 4, 9]))
+  b = net.add_node(np.ones([9, 4, 5]))
+  net.connect(a[2], b[1])
+  net.connect(a[3], b[0])
+  c = net.batched_contract_between(a, b, a[0], b[2])
+  np.testing.assert_allclose(c.tensor, 36 * np.ones([5, 3]))
+  
+def test_pairwise_reduction():
+  net = batchtensornetwork.BatchTensorNetwork(backend="tensorflow")
+  a = net.add_node(np.ones([10, 2, 2]))
+  b = batchtensornetwork.pairwise_reduction(net, a, a[0])
+  np.testing.assert_allclose(b.tensor, 512 * np.ones([2, 2]))
+  
+def test_pairwise_reduction_multiple_edges():
+  net = batchtensornetwork.BatchTensorNetwork(backend="tensorflow")
+  a = net.add_node(np.ones([3, 5, 6, 2, 2]))
+  b = batchtensornetwork.pairwise_reduction(net, a, a[1])
+  np.testing.assert_allclose(b.tensor, 16 * np.ones([3, 6, 2, 2]))
diff --git a/experiments/MPS_classifier/classifier_test.py b/experiments/MPS_classifier/classifier_test.py
--- a/experiments/MPS_classifier/classifier_test.py
+++ b/experiments/MPS_classifier/classifier_test.py
@@ -1,74 +1,73 @@
-# Copyright 2019 The TensorNetwork Authors
-#
-# Licensed under the Apache License, Version 2.0 (the "License");
-# you may not use this file except in compliance with the License.
-# You may obtain a copy of the License at
-#
-#      http://www.apache.org/licenses/LICENSE-2.0
-#
-# Unless required by applicable law or agreed to in writing, software
-# distributed under the License is distributed on an "AS IS" BASIS,
-# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
-# See the License for the specific language governing permissions and
-# limitations under the License.
-
-from __future__ import absolute_import
-from __future__ import division
-from __future__ import print_function
-
-import numpy as np
-import tensorflow as tf
-from experiments.MPS_classifier import classifier
-tf.enable_v2_behavior()
-
-
-def test_random_initializer_shapes():
-  x = classifier.random_initializer(2, 3)
-  assert x.shape == (2, 3, 3)
-  y = classifier.random_initializer(2, 3, boundary=True)
-  assert y.shape == (2, 3)
-
-def test_environment_shapes():
-  env = classifier.Environment(5, 2, 7)
-  assert list(env.vector.shape) == [2, 7]
-  assert list(env.matrices.shape) == [4, 2, 7, 7]
-
-def test_environment_network():
-  env = classifier.Environment(5, 2, 7)
-  data = tf.cast(np.random.random([20, 5, 2]), dtype=tf.float32)
-  net, var_nodes, data_nodes = env.create_network(data[:, 1:], data[:, 0])
-  net.check_correct()
-  np.testing.assert_allclose(data_nodes[0].tensor.numpy(), data[:, 0])
-  np.testing.assert_allclose(data_nodes[1].tensor.numpy(), data[:, 1:])
-
-def test_environment_predict():
-  env = classifier.Environment(5, 2, 7)
-  data = tf.cast(np.random.random([20, 5, 2]), dtype=tf.float32)
-  # Compare with serialized contraction
-  pred = tf.einsum("bs,sr->br", data[:, 0], env.vector)
-  for i in range(4):
-    pred = tf.einsum("bl,bs,slr->br", pred, data[:, i + 1], env.matrices[i])
-  np.testing.assert_allclose(env.predict(data[:, 1:], data[:, 0]).numpy(),
-                             pred.numpy(), atol=1e-5)
-
-def test_mps_shapes():
-  mps = classifier.MatrixProductState(11, 10, 2, 7)
-  assert list(mps.tensors[0].shape) == [2, 7]
-  assert list(mps.tensors[1].shape) == [4, 2,  7, 7]
-  assert list(mps.tensors[2].shape) == [10, 7, 7]
-  assert list(mps.tensors[3].shape) == [4, 2, 7, 7]
-  assert list(mps.tensors[4].shape) == [2, 7]
-
-def test_calculate_flx():
-  mps = classifier.MatrixProductState(11, 10, 2, 7)
-  data = tf.cast(np.random.random([20, 10, 2]), dtype=mps.dtype)
-  # Compare with serialized contraction
-  flx_left = tf.einsum("bs,sr->br", data[:, 0], mps.tensors[0])
-  flx_right = tf.einsum("bs,sr->br", data[:, -1], mps.tensors[4])
-  for i in range(4):
-    flx_left = tf.einsum("bl,bs,slr->br", flx_left, data[:, i + 1],
-                         mps.tensors[1][i])
-    flx_right = tf.einsum("bl, bs,slr->br", flx_right, data[:, 8 - i],
-                          mps.tensors[3][i])
-  flx = tf.einsum("bl,olr,br->bo", flx_left, mps.tensors[2], flx_right)
-  np.testing.assert_allclose(flx.numpy(), mps.flx(data).numpy(), atol=1e-5)
+# Copyright 2019 The TensorNetwork Authors
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#      http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+from __future__ import absolute_import
+from __future__ import division
+from __future__ import print_function
+
+import numpy as np
+import tensorflow as tf
+from experiments.MPS_classifier import classifier
+
+
+def test_random_initializer_shapes():
+  x = classifier.random_initializer(2, 3)
+  assert x.shape == (2, 3, 3)
+  y = classifier.random_initializer(2, 3, boundary=True)
+  assert y.shape == (2, 3)
+
+def test_environment_shapes():
+  env = classifier.Environment(5, 2, 7)
+  assert list(env.vector.shape) == [2, 7]
+  assert list(env.matrices.shape) == [4, 2, 7, 7]
+
+def test_environment_network():
+  env = classifier.Environment(5, 2, 7)
+  data = tf.cast(np.random.random([20, 5, 2]), dtype=tf.float32)
+  net, var_nodes, data_nodes = env.create_network(data[:, 1:], data[:, 0])
+  net.check_correct()
+  np.testing.assert_allclose(data_nodes[0].tensor.numpy(), data[:, 0])
+  np.testing.assert_allclose(data_nodes[1].tensor.numpy(), data[:, 1:])
+
+def test_environment_predict():
+  env = classifier.Environment(5, 2, 7)
+  data = tf.cast(np.random.random([20, 5, 2]), dtype=tf.float32)
+  # Compare with serialized contraction
+  pred = tf.einsum("bs,sr->br", data[:, 0], env.vector)
+  for i in range(4):
+    pred = tf.einsum("bl,bs,slr->br", pred, data[:, i + 1], env.matrices[i])
+  np.testing.assert_allclose(env.predict(data[:, 1:], data[:, 0]).numpy(),
+                             pred.numpy(), atol=1e-5)
+
+def test_mps_shapes():
+  mps = classifier.MatrixProductState(11, 10, 2, 7)
+  assert list(mps.tensors[0].shape) == [2, 7]
+  assert list(mps.tensors[1].shape) == [4, 2,  7, 7]
+  assert list(mps.tensors[2].shape) == [10, 7, 7]
+  assert list(mps.tensors[3].shape) == [4, 2, 7, 7]
+  assert list(mps.tensors[4].shape) == [2, 7]
+
+def test_calculate_flx():
+  mps = classifier.MatrixProductState(11, 10, 2, 7)
+  data = tf.cast(np.random.random([20, 10, 2]), dtype=mps.dtype)
+  # Compare with serialized contraction
+  flx_left = tf.einsum("bs,sr->br", data[:, 0], mps.tensors[0])
+  flx_right = tf.einsum("bs,sr->br", data[:, -1], mps.tensors[4])
+  for i in range(4):
+    flx_left = tf.einsum("bl,bs,slr->br", flx_left, data[:, i + 1],
+                         mps.tensors[1][i])
+    flx_right = tf.einsum("bl, bs,slr->br", flx_right, data[:, 8 - i],
+                          mps.tensors[3][i])
+  flx = tf.einsum("bl,olr,br->bo", flx_left, mps.tensors[2], flx_right)
+  np.testing.assert_allclose(flx.numpy(), mps.flx(data).numpy(), atol=1e-5)
diff --git a/experiments/tree_tensor_network/ttn_test.py b/experiments/tree_tensor_network/ttn_test.py
--- a/experiments/tree_tensor_network/ttn_test.py
+++ b/experiments/tree_tensor_network/ttn_test.py
@@ -17,14 +17,9 @@
 import pytest
 import numpy as np
 import tensorflow as tf
-from jax.config import config
-config.update("jax_enable_x64", True)
-import tensornetwork as tn
 import examples.wavefunctions.wavefunctions as wf
 import experiments.tree_tensor_network as ttn
 
-tf.enable_v2_behavior()
-
 
 def test_opt(backend):
   if backend == "tensorflow":
diff --git a/tensornetwork/backends/jax/jax_backend_test.py b/tensornetwork/backends/jax/jax_backend_test.py
--- a/tensornetwork/backends/jax/jax_backend_test.py
+++ b/tensornetwork/backends/jax/jax_backend_test.py
@@ -7,7 +7,6 @@
 import jax
 import pytest
 from tensornetwork.backends.jax import jax_backend
-import tensornetwork.config as config_file
 
 np_randn_dtypes = [np.float32, np.float16, np.float64]
 np_dtypes = np_randn_dtypes + [np.complex64, np.complex128]
diff --git a/tensornetwork/backends/numpy/numpy_backend_test.py b/tensornetwork/backends/numpy/numpy_backend_test.py
--- a/tensornetwork/backends/numpy/numpy_backend_test.py
+++ b/tensornetwork/backends/numpy/numpy_backend_test.py
@@ -6,7 +6,7 @@
 import numpy as np
 import pytest
 from tensornetwork.backends.numpy import numpy_backend
-import tensornetwork.config as config_file
+
 
 np_randn_dtypes = [np.float32, np.float16, np.float64]
 np_dtypes = np_randn_dtypes + [np.complex64, np.complex128]
diff --git a/tensornetwork/backends/pytorch/pytorch_backend_test.py b/tensornetwork/backends/pytorch/pytorch_backend_test.py
--- a/tensornetwork/backends/pytorch/pytorch_backend_test.py
+++ b/tensornetwork/backends/pytorch/pytorch_backend_test.py
@@ -6,7 +6,7 @@
 from tensornetwork.backends.pytorch import pytorch_backend
 import torch
 import pytest
-import tensornetwork.config as config_file
+
 
 torch_dtypes = [torch.float32, torch.float64, torch.int32]
 torch_eye_dtypes = [torch.float32, torch.float64, torch.int32, torch.int64]
diff --git a/tensornetwork/backends/tensorflow/tensordot2_test.py b/tensornetwork/backends/tensorflow/tensordot2_test.py
--- a/tensornetwork/backends/tensorflow/tensordot2_test.py
+++ b/tensornetwork/backends/tensorflow/tensordot2_test.py
@@ -37,7 +37,7 @@
 from tensornetwork.backends.tensorflow import tensordot2
 import pytest
 
-tf.compat.v1.enable_v2_behavior()
+
 _MAXDIM = 5
 
 
diff --git a/tensornetwork/backends/tensorflow/tensorflow_backend_test.py b/tensornetwork/backends/tensorflow/tensorflow_backend_test.py
--- a/tensornetwork/backends/tensorflow/tensorflow_backend_test.py
+++ b/tensornetwork/backends/tensorflow/tensorflow_backend_test.py
@@ -5,9 +5,8 @@
 import numpy as np
 import tensorflow as tf
 import pytest
-import tensornetwork.config as config_file
 from tensornetwork.backends.tensorflow import tensorflow_backend
-tf.compat.v1.enable_v2_behavior()
+
 
 tf_randn_dtypes = [tf.float32, tf.float16, tf.float64]
 tf_dtypes = tf_randn_dtypes + [tf.complex128, tf.complex64]
diff --git a/tensornetwork/backends/tensorflow/tensorflow_tensornetwork_test.py b/tensornetwork/backends/tensorflow/tensorflow_tensornetwork_test.py
--- a/tensornetwork/backends/tensorflow/tensorflow_tensornetwork_test.py
+++ b/tensornetwork/backends/tensorflow/tensorflow_tensornetwork_test.py
@@ -4,7 +4,6 @@
 from __future__ import print_function
 import numpy as np
 import tensorflow as tf
-tf.compat.v1.enable_v2_behavior()
 import tensornetwork
 
 
diff --git a/tensornetwork/contractors/cost_calculators_test.py b/tensornetwork/contractors/cost_calculators_test.py
--- a/tensornetwork/contractors/cost_calculators_test.py
+++ b/tensornetwork/contractors/cost_calculators_test.py
@@ -18,7 +18,6 @@
 from __future__ import print_function
 import numpy as np
 import pytest
-from typing import List, Optional, Tuple
 from tensornetwork.contractors import cost_calculators
 from tensornetwork import network
 
diff --git a/tensornetwork/contractors/greedy_contractor_test.py b/tensornetwork/contractors/greedy_contractor_test.py
--- a/tensornetwork/contractors/greedy_contractor_test.py
+++ b/tensornetwork/contractors/greedy_contractor_test.py
@@ -17,8 +17,6 @@
 from __future__ import division
 from __future__ import print_function
 import numpy as np
-import pytest
-from typing import List, Optional, Tuple
 from tensornetwork.contractors import greedy_contractor
 from tensornetwork import network
 
diff --git a/tensornetwork/contractors/naive_contractor_test.py b/tensornetwork/contractors/naive_contractor_test.py
--- a/tensornetwork/contractors/naive_contractor_test.py
+++ b/tensornetwork/contractors/naive_contractor_test.py
@@ -22,8 +22,6 @@
 
 from tensornetwork import network
 from tensornetwork.contractors import naive_contractor
-import tensorflow as tf
-tf.enable_v2_behavior()
 naive = naive_contractor.naive
 
 
diff --git a/tensornetwork/contractors/opt_einsum_paths/path_contractors_test.py b/tensornetwork/contractors/opt_einsum_paths/path_contractors_test.py
--- a/tensornetwork/contractors/opt_einsum_paths/path_contractors_test.py
+++ b/tensornetwork/contractors/opt_einsum_paths/path_contractors_test.py
@@ -16,8 +16,6 @@
 import pytest
 import tensornetwork
 from tensornetwork.contractors.opt_einsum_paths import path_contractors
-import tensorflow as tf
-tf.enable_v2_behavior()
 
 
 @pytest.fixture(name="path_algorithm",
diff --git a/tensornetwork/contractors/stochastic_contractor_test.py b/tensornetwork/contractors/stochastic_contractor_test.py
--- a/tensornetwork/contractors/stochastic_contractor_test.py
+++ b/tensornetwork/contractors/stochastic_contractor_test.py
@@ -21,7 +21,6 @@
 import tensorflow as tf
 from tensornetwork import network
 from tensornetwork.contractors import stochastic_contractor
-tf.compat.v1.enable_v2_behavior()
 
 
 class StochasticTest(tf.test.TestCase):
diff --git a/tensornetwork/ncon_interface_test.py b/tensornetwork/ncon_interface_test.py
--- a/tensornetwork/ncon_interface_test.py
+++ b/tensornetwork/ncon_interface_test.py
@@ -17,128 +17,134 @@
 from __future__ import print_function
 import pytest
 import numpy as np
-import tensorflow as tf
-tf.compat.v1.enable_v2_behavior()
 from tensornetwork import ncon_interface
 from tensornetwork.contractors.naive_contractor import naive
 
 
-def test_sanity_check():
-  result = ncon_interface.ncon([tf.ones(
-      (2, 2)), tf.ones((2, 2))], [(-1, 1), (1, -2)])
-  np.testing.assert_allclose(result, tf.ones((2, 2)) * 2)
+def test_sanity_check(backend):
+  result = ncon_interface.ncon([np.ones(
+      (2, 2)), np.ones((2, 2))], [(-1, 1), (1, -2)], backend=backend)
+  np.testing.assert_allclose(result, np.ones((2, 2)) * 2)
 
 
-def test_order_spec():
-  a = tf.ones((2, 2))
+def test_order_spec(backend):
+  a = np.ones((2, 2))
 
-  result = ncon_interface.ncon([a, a], [(-1, 1), (1, -2)], out_order=[-1, -2])
-  np.testing.assert_allclose(result, tf.ones((2, 2)) * 2)
+  result = ncon_interface.ncon(
+      [a, a], [(-1, 1), (1, -2)], out_order=[-1, -2], backend=backend)
+  np.testing.assert_allclose(result, np.ones((2, 2)) * 2)
 
-  result = ncon_interface.ncon([a, a], [(-1, 1), (1, -2)], con_order=[1])
-  np.testing.assert_allclose(result, tf.ones((2, 2)) * 2)
+  result = ncon_interface.ncon(
+      [a, a], [(-1, 1), (1, -2)], con_order=[1], backend=backend)
+  np.testing.assert_allclose(result, np.ones((2, 2)) * 2)
 
   result = ncon_interface.ncon([a, a], [(-1, 1), (1, -2)],
                                con_order=[1],
                                out_order=[-1, -2])
-  np.testing.assert_allclose(result, tf.ones((2, 2)) * 2)
+  np.testing.assert_allclose(result, np.ones((2, 2)) * 2)
 
 
-def test_order_spec_noninteger():
-  a = tf.ones((2, 2))
+def test_order_spec_noninteger(backend):
+  a = np.ones((2, 2))
   result = ncon_interface.ncon([a, a], [('o1', 'i'), ('i', 'o2')],
                                con_order=['i'],
-                               out_order=['o1', 'o2'])
-  np.testing.assert_allclose(result, tf.ones((2, 2)) * 2)
+                               out_order=['o1', 'o2'], backend=backend)
+  np.testing.assert_allclose(result, np.ones((2, 2)) * 2)
 
 
-def test_invalid_network():
-  a = tf.ones((2, 2))
+def test_invalid_network(backend):
+  a = np.ones((2, 2))
   with pytest.raises(ValueError):
-    ncon_interface.ncon([a, a], [(1, 2), (2, 1), (1, 2)])
+    ncon_interface.ncon([a, a], [(1, 2), (2, 1), (1, 2)], backend=backend)
   with pytest.raises(ValueError):
-    ncon_interface.ncon([a, a], [(1, 2), (2, 2)])
+    ncon_interface.ncon([a, a], [(1, 2), (2, 2)], backend=backend)
   with pytest.raises(ValueError):
-    ncon_interface.ncon([a, a], [(1, 2), (3, 1)])
+    ncon_interface.ncon([a, a], [(1, 2), (3, 1)], backend=backend)
   with pytest.raises(ValueError):
-    ncon_interface.ncon([a, a], [(1, 2), (2, 0.1)])
+    ncon_interface.ncon([a, a], [(1, 2), (2, 0.1)], backend=backend)
   with pytest.raises(ValueError):
-    ncon_interface.ncon([a, a], [(1, 2), (2, 't')])
+    ncon_interface.ncon([a, a], [(1, 2), (2, 't')], backend=backend)
   with pytest.raises(ValueError):
-    ncon_interface.ncon([a, a], [(0, 1), (1, 0)])
+    ncon_interface.ncon([a, a], [(0, 1), (1, 0)], backend=backend)
   with pytest.raises(ValueError):
-    ncon_interface.ncon([a, a], [(1,), (1, 2)])
+    ncon_interface.ncon([a, a], [(1,), (1, 2)], backend=backend)
 
 
-def test_invalid_order():
-  a = tf.ones((2, 2))
+def test_invalid_order(backend):
+  a = np.ones((2, 2))
   with pytest.raises(ValueError):
-    ncon_interface.ncon([a, a], [(1, 2), (2, 1)], con_order=[2, 3])
+    ncon_interface.ncon(
+        [a, a], [(1, 2), (2, 1)], con_order=[2, 3], backend=backend)
   with pytest.raises(ValueError):
-    ncon_interface.ncon([a, a], [(1, 2), (2, 1)], out_order=[-1])
+    ncon_interface.ncon(
+        [a, a], [(1, 2), (2, 1)], out_order=[-1], backend=backend)
   with pytest.raises(ValueError):
     ncon_interface.ncon([a, a], [('i1', 'i2'), ('i1', 'i2')],
                         con_order=['i1'],
-                        out_order=[])
+                        out_order=[], backend=backend)
   with pytest.raises(ValueError):
     ncon_interface.ncon([a, a], [('i1', 'i2'), ('i1', 'i2')],
                         con_order=['i1', 'i2'],
-                        out_order=['i1'])
+                        out_order=['i1'], backend=backend)
   with pytest.raises(ValueError):
     ncon_interface.ncon([a, a], [('i1', 'i2'), ('i1', 'i2')],
                         con_order=['i1', 'i1', 'i2'],
-                        out_order=[])
+                        out_order=[], backend=backend)
 
 
-def test_out_of_order_contraction():
-  a = tf.ones((2, 2, 2))
+def test_out_of_order_contraction(backend):
+  a = np.ones((2, 2, 2))
   with pytest.warns(UserWarning, match='Suboptimal ordering'):
-    ncon_interface.ncon([a, a, a], [(-1, 1, 3), (1, 3, 2), (2, -2, -3)])
+    ncon_interface.ncon(
+        [a, a, a], [(-1, 1, 3), (1, 3, 2), (2, -2, -3)], backend=backend)
 
 
-def test_output_order():
+def test_output_order(backend):
   a = np.random.randn(2, 2)
-  res = ncon_interface.ncon([a], [(-2, -1)])
+  res = ncon_interface.ncon([a], [(-2, -1)], backend=backend)
   np.testing.assert_allclose(res, a.transpose())
 
 
-def test_outer_product():
+def test_outer_product(backend):
+  if backend == "jax":
+    pytest.skip("Jax outer product support is currently broken.")
   a = np.array([1, 2, 3])
   b = np.array([1, 2])
-  res = ncon_interface.ncon([a, b], [(-1,), (-2,)])
+  res = ncon_interface.ncon([a, b], [(-1,), (-2,)], backend=backend)
   np.testing.assert_allclose(res, np.kron(a, b).reshape((3, 2)))
-  res = ncon_interface.ncon([a, a, a, a], [(1,), (1,), (2,), (2,)])
-  assert res.numpy() == 196
+  res = ncon_interface.ncon(
+      [a, a, a, a], [(1,), (1,), (2,), (2,)], backend=backend)
+  np.testing.assert_allclose(res, 196)
 
 
-def test_trace():
-  a = tf.ones((2, 2))
-  res = ncon_interface.ncon([a], [(1, 1)])
-  assert res.numpy() == 2
+def test_trace(backend):
+  a = np.ones((2, 2))
+  res = ncon_interface.ncon([a], [(1, 1)], backend=backend)
+  np.testing.assert_allclose(res, 2)
 
 
-def test_small_matmul():
+def test_small_matmul(backend):
   a = np.random.randn(2, 2)
   b = np.random.randn(2, 2)
-  res = ncon_interface.ncon([a, b], [(1, -1), (1, -2)])
+  res = ncon_interface.ncon(
+      [a, b], [(1, -1), (1, -2)], backend=backend)
   np.testing.assert_allclose(res, a.transpose() @ b)
 
 
-def test_contraction():
+def test_contraction(backend):
   a = np.random.randn(2, 2, 2)
-  res = ncon_interface.ncon([a, a, a], [(-1, 1, 2), (1, 2, 3), (3, -2, -3)])
+  res = ncon_interface.ncon(
+      [a, a, a], [(-1, 1, 2), (1, 2, 3), (3, -2, -3)], backend=backend)
   res_np = a.reshape((2, 4)) @ a.reshape((4, 2)) @ a.reshape((2, 4))
   res_np = res_np.reshape((2, 2, 2))
   np.testing.assert_allclose(res, res_np)
 
 
-def test_backend_network():
+def test_backend_network(backend):
   a = np.random.randn(2, 2, 2)
   tn, _, _ = ncon_interface.ncon_network([a, a, a], [(-1, 1, 2), (1, 2, 3),
                                                      (3, -2, -3)],
-                                         backend="numpy")
-  assert tn.backend.name == "numpy"
-
+                                         backend=backend)
   res = naive(tn).get_final_node().get_tensor()
   res_np = a.reshape((2, 4)) @ a.reshape((4, 2)) @ a.reshape((2, 4))
   res_np = res_np.reshape((2, 2, 2))
diff --git a/tensornetwork/network_components_test.py b/tensornetwork/network_components_test.py
--- a/tensornetwork/network_components_test.py
+++ b/tensornetwork/network_components_test.py
@@ -6,7 +6,6 @@
 from tensornetwork.network_components import Node, CopyNode, Edge
 import tensornetwork
 
-tf.compat.v1.enable_v2_behavior()
 
 string_type = h5py.special_dtype(vlen=str)
 
diff --git a/tensornetwork/tests/network_test.py b/tensornetwork/tests/network_test.py
--- a/tensornetwork/tests/network_test.py
+++ b/tensornetwork/tests/network_test.py
@@ -22,12 +22,8 @@
 import torch
 import jax
 from jax.config import config
-import tensornetwork.config as config_file
 
 
-config.update("jax_enable_x64", True)
-tf.compat.v1.enable_v2_behavior()
-
 np_dtypes = [np.float32, np.float64, np.complex64, np.complex128, np.int32]
 tf_dtypes = [tf.float32, tf.float64, tf.complex64, tf.complex128, tf.int32]
 torch_dtypes = [torch.float32, torch.float64, torch.int32, torch.int64]
diff --git a/tensornetwork/tests/serialization_test.py b/tensornetwork/tests/serialization_test.py
--- a/tensornetwork/tests/serialization_test.py
+++ b/tensornetwork/tests/serialization_test.py
@@ -18,15 +18,9 @@
 import tensornetwork
 import io
 import numpy as np
-import tensorflow as tf
-from jax.config import config
 import h5py
 
 
-config.update("jax_enable_x64", True)
-tf.compat.v1.enable_v2_behavior()
-
-
 def test_save_makes_hdf5_file(tmp_path, backend):
   net = tensornetwork.TensorNetwork(backend=backend)
   a = net.add_node(np.ones((2, 2, 2)))
diff --git a/tensornetwork/tests/split_node_test.py b/tensornetwork/tests/split_node_test.py
--- a/tensornetwork/tests/split_node_test.py
+++ b/tensornetwork/tests/split_node_test.py
@@ -18,13 +18,6 @@
 import tensornetwork
 import pytest
 import numpy as np
-import tensorflow as tf
-from jax.config import config
-import tensornetwork.config as config_file
-
-
-config.update("jax_enable_x64", True)
-tf.compat.v1.enable_v2_behavior()
 
 
 def test_split_node_qr_disable(backend):
diff --git a/tensornetwork/tests/tensornetwork_test.py b/tensornetwork/tests/tensornetwork_test.py
--- a/tensornetwork/tests/tensornetwork_test.py
+++ b/tensornetwork/tests/tensornetwork_test.py
@@ -21,13 +21,8 @@
 import tensorflow as tf
 import torch
 import jax
-from jax.config import config
-import tensornetwork.config as config_file
 
 
-config.update("jax_enable_x64", True)
-tf.compat.v1.enable_v2_behavior()
-
 np_dtypes = [np.float32, np.float64, np.complex64, np.complex128, np.int32]
 tf_dtypes = [tf.float32, tf.float64, tf.complex64, tf.complex128, tf.int32]
 torch_dtypes = [torch.float32, torch.float64, torch.int32, torch.int64]