Use opt_einsum to optimize contraction orders

Author: JanLuca

varipeps/contractions/apply.py

@@ -35,16 +35,18 @@ def __class_getitem__(cls, name: str) -> Checkpointing_Cache:
 _ncon_jitted = jax.jit(_jittable_ncon, static_argnums=(1, 2, 3, 4, 5), inline=True)
 
 
+@partial(
+    jax.jit, static_argnames=("name", "disable_identity_check", "custom_definition")
+)
 def apply_contraction(
     name: str,
     peps_tensors: Sequence[jnp.ndarray],
     peps_tensor_objs: Sequence[PEPS_Tensor],
     additional_tensors: Sequence[jnp.ndarray],
     *,
-    disable_identity_check: bool = False,
+    disable_identity_check: bool = True,
     custom_definition: Optional[Definition] = None,
     config: VariPEPS_Config = varipeps_config,
-    _jitable: bool = False,
 ) -> jnp.ndarray:
     """
     Apply a contraction to a list of tensors.
@@ -83,8 +85,7 @@ class definition for details.
         )
 
     if (
-        not _jitable
-        and not disable_identity_check
+        not disable_identity_check
         and not all(isinstance(t, jax.core.Tracer) for t in peps_tensors)
         and not all(isinstance(to.tensor, jax.core.Tracer) for to in peps_tensor_objs)
         and not all(
@@ -133,36 +134,19 @@ class definition for details.
 
     tensors += additional_tensors
 
-    if _jitable:
-        if config.checkpointing_ncon:
-            f = jax.checkpoint(_ncon_jitted, static_argnums=(1, 2, 3, 4, 5))
-        else:
-            f = _ncon_jitted
-
-        return f(
-            tensors,
-            contraction["ncon_flat_network"],
-            contraction["ncon_sizes"],
-            contraction["ncon_con_order"],
-            contraction["ncon_out_order"],
-            tn.backends.backend_factory.get_backend("jax"),
-        )
+    tensor_shapes = tuple(tuple(e.shape) for e in tensors)
 
-    if config.checkpointing_ncon:
-        f = jax.checkpoint(
-            partial(
-                tn.ncon, network_structure=contraction["ncon_network"], backend="jax"
-            )
-        )
-    else:
-        f = partial(
-            tn.ncon, network_structure=contraction["ncon_network"], backend="jax"
+    path = contraction["einsum_cache"].get(tensor_shapes)
+
+    if path is None:
+        path, _ = jnp.einsum_path(
+            contraction["einsum_network"],
+            *tensors,
+            optimize="optimal" if len(tensors) < 10 else "dp",
         )
+        contraction["einsum_cache"][tensor_shapes] = path
 
-    return f(tensors)
+    return jnp.einsum(contraction["einsum_network"], *tensors, optimize=path)
 
 
-apply_contraction_jitted = jax.jit(
-    partial(apply_contraction, _jitable=True),
-    static_argnames=("name", "disable_identity_check", "custom_definition"),
-)
+apply_contraction_jitted = apply_contraction

varipeps/contractions/definitions.py

@@ -4,6 +4,7 @@
 
 from collections import Counter
 
+import opt_einsum
 import tensornetwork as tn
 
 from typing import Dict, Optional, Union, List, Tuple, Sequence
@@ -116,6 +117,31 @@ def _create_filter_and_network(
             network_additional_tensors,
         )
 
+    @classmethod
+    def _convert_to_einsum(cls, network):
+        max_contracted = 0
+        min_open = 0
+
+        result = tuple([] for _ in range(len(network)))
+
+        for ti, t in enumerate(network):
+            for i in t:
+                if i < 0 and i < min_open:
+                    min_open = i
+                elif i > 0 and i > max_contracted:
+                    max_contracted = i
+                if max_contracted >= (min_open + 52):
+                    raise ValueError("Letters in conversion are overlapping.")
+
+                result[ti].append(opt_einsum.get_symbol(i))
+
+        result = ["".join(e) for e in result]
+        open_result = [opt_einsum.get_symbol(i) for i in range(-1, min_open - 1, -1)]
+
+        result = f"{','.join(result)}->{''.join(open_result)}"
+
+        return result
+
     @classmethod
     def _process_def(cls, e, name):
         (
@@ -129,6 +155,8 @@ def _process_def(cls, e, name):
             j for i in network_peps_tensors for j in i
         ] + network_additional_tensors
 
+        einsum_network = cls._convert_to_einsum(ncon_network)
+
         flatted_ncon_list = [j for i in ncon_network for j in i]
         counter_ncon_list = Counter(flatted_ncon_list)
         for ind, c in counter_ncon_list.items():
@@ -142,28 +170,13 @@ def _process_def(cls, e, name):
                 f'Non-monotonous indices in definition "{name}". Please check!'
             )
 
-        (
-            mapped_ncon_network,
-            mapping,
-        ) = tn.ncon_interface._canonicalize_network_structure(ncon_network)
-        flat_network = tuple(l for sublist in mapped_ncon_network for l in sublist)
-        unique_flat_network = list(set(flat_network))
-
-        out_order = tuple(
-            sorted([l for l in unique_flat_network if l < 0], reverse=True)
-        )
-        con_order = tuple(sorted([l for l in unique_flat_network if l > 0]))
-        sizes = tuple(len(l) for l in ncon_network)
-
         e["filter_peps_tensors"] = filter_peps_tensors
         e["filter_additional_tensors"] = filter_additional_tensors
         e["network_peps_tensors"] = network_peps_tensors
         e["network_additional_tensors"] = network_additional_tensors
         e["ncon_network"] = ncon_network
-        e["ncon_flat_network"] = flat_network
-        e["ncon_sizes"] = sizes
-        e["ncon_con_order"] = con_order
-        e["ncon_out_order"] = out_order
+        e["einsum_network"] = einsum_network
+        e["einsum_cache"] = dict()
 
     @classmethod
     def _prepare_defs(cls):