Repository: google/TensorNetwork

Issue:
ncon should take a BaseBackend object
Since we are now supporting passing `BaseBackend` to `Node` constructors, we should also
allow it in `ncon`


Discussion:
Yes. Should be easy to fix right? I think we just need to update the typing
that's right, I think it should be woring already
has anyone taken the lead on this?
No, but this should be pretty easy for a contributor to do. 
Hi I would like to work on this issue. If I understand this issue correctly, the function parameters for ncon and ncon_network should be updated to accept both Text and BaseBackend object?

Resolving patch:
diff --git a/tensornetwork/ncon_interface.py b/tensornetwork/ncon_interface.py
--- a/tensornetwork/ncon_interface.py
+++ b/tensornetwork/ncon_interface.py
@@ -18,6 +18,8 @@
 from tensornetwork import network_components
 from tensornetwork.backend_contextmanager import get_default_backend
 from tensornetwork.backends import backend_factory
+from tensornetwork.backends.base_backend import BaseBackend
+
 Tensor = Any
 
 
@@ -26,7 +28,7 @@ def ncon(
     network_structure: Sequence[Sequence],
     con_order: Optional[Sequence] = None,
     out_order: Optional[Sequence] = None,
-    backend: Optional[Text] = None
+    backend: Optional[Union[Text, BaseBackend]] = None
 ) -> Union[network_components.BaseNode, Tensor]:
   r"""Contracts a list of tensors or nodes according to a tensor network 
     specification.
@@ -79,16 +81,18 @@ def ncon(
       if all elements of `tensors` are `BaseNode` objects, else
       it is returned as a `Tensor` object.
     """
-  if backend and (backend not in backend_factory._BACKENDS):
-    raise ValueError("Backend '{}' does not exist".format(backend))
   if backend is None:
     backend = get_default_backend()
+  if isinstance(backend, BaseBackend):
+    backend_obj = backend
+  else:
+    backend_obj = backend_factory.get_backend(backend)
 
   are_nodes = [isinstance(t, network_components.BaseNode) for t in tensors]
   nodes = {t for t in tensors if isinstance(t, network_components.BaseNode)}
-  if not all([n.backend.name == backend for n in nodes]):
-    raise ValueError(
-        "Some nodes have backends different from '{}'".format(backend))
+  if not all([n.backend.name == backend_obj.name for n in nodes]):
+    raise ValueError("Some nodes have backends different from '{}'".format(
+        backend_obj.name))
 
   _tensors = []
   for t in tensors:
@@ -102,7 +106,7 @@ def ncon(
       network_structure,
       con_order=con_order,
       out_order=out_order,
-      backend=backend)
+      backend=backend_obj)
 
   nodes = set(nodes)  # we don't need the ordering here
 
@@ -160,7 +164,7 @@ def ncon_network(
     network_structure: Sequence[Sequence],
     con_order: Optional[Sequence] = None,
     out_order: Optional[Sequence] = None,
-    backend: Optional[Text] = None
+    backend: Optional[Union[Text, BaseBackend]] = None
 ) -> Tuple[List[network_components.BaseNode], List[network_components.Edge],
            List[network_components.Edge]]:
   r"""Creates a network from a list of tensors according to `tensors`.
@@ -182,9 +186,8 @@ def ncon_network(
       network_structure: List of lists specifying the tensor network.
       con_order: List of edge labels specifying the contraction order.
       out_order: List of edge labels specifying the output order.
-      backend: String specifying the backend to use. Defaults to the default
-        TensorNetwork backend.
-
+      backend: String or BaseBackend object specifying the backend to use. 
+        Defaults to the default TensorNetwork backend.
     Returns:
       nodes: List of constructed nodes in the same order as given in `tensors`.
       con_edges: List of internal `Edge` objects in contraction order.
@@ -247,8 +250,9 @@ def ncon_network(
 
 
 def _build_network(
-    tensors: Sequence[Tensor], network_structure: Sequence[Sequence],
-    backend: Text
+    tensors: Sequence[Tensor],
+    network_structure: Sequence[Sequence],
+    backend: Optional[Union[BaseBackend, Text]] = None,
 ) -> Tuple[List[network_components.BaseNode], Dict[Any,
                                                    network_components.Edge]]:
   nodes = []


Test patch:
diff --git a/tensornetwork/tests/ncon_interface_test.py b/tensornetwork/tests/ncon_interface_test.py
--- a/tensornetwork/tests/ncon_interface_test.py
+++ b/tensornetwork/tests/ncon_interface_test.py
@@ -15,9 +15,23 @@
 import numpy as np
 from tensornetwork import BaseNode, Node
 from tensornetwork import ncon_interface
+from tensornetwork.backends.backend_factory import get_backend
 from tensornetwork.contractors import greedy
 
 
+@pytest.fixture(
+    name="backend",
+    params=[
+        'numpy',
+        get_backend('numpy'), 'jax',
+        get_backend('jax'), 'pytorch',
+        get_backend('pytorch'), 'tensorflow',
+        get_backend('tensorflow')
+    ])
+def backends(request):
+  return request.param
+
+
 def test_sanity_check(backend):
   t1, t2 = np.ones((2, 2)), np.ones((2, 2))
   result_1 = ncon_interface.ncon([t1, t2], [(-1, 1), (1, -2)], backend=backend)