Repository: dgasmith/opt_einsum

Issue:
Remove 'remaining' from PathInfo? (incurs n^2 memory overhead)
When building the ``contraction_list``, ``opt_einsum`` handles a list of 'remaining'  terms, and at each new step these are copied (``input_list`` below):

https://github.com/dgasmith/opt_einsum/blob/2cbd28e612ff0af54656c34d7f6c239efe524b12/opt_einsum/contract.py#L306

A single, consumed version of this list is needed for constructing the items in the contraction list. However all the copies at each stage are not needed in the actual contraction (``_core_contract``) function - they are just for printing in the ``PathInfo`` object. 

The problem is that each copied list of remaining terms is of average length `n / 2` and there are generally `n - 1` of them, incurring a `n^2` memory overhead, which caused problems for us (10s of gigabytes of problems!) in some large contractions. Here's an extreme sample of such a gigantic ``PathInfo``:

![Screenshot from 2020-07-15 19-11-55](https://user-images.githubusercontent.com/8982598/87640894-486a1400-c73f-11ea-98a1-dab04d3e02f8.png)

With 50 contractions each of several thousand tensors there was a total of nearly 100,000,000 'remaining' lists of strings.

I'd suggest either:
* a) *Easiest* - we just remove the 'remaining' terms from the ``PathInfo`` object, I'd be fine with this but possibly its useful information for others?
* b) *Medium* - the 'remaining' terms are only accumulated below a certain threshold size
* c) *Hardest* - they are generated lazily (and maybe optionally) by the ``PathInfo`` object upon request only (e.g. printing)

Discussion:
codecov[bot]: # [Codecov](https://codecov.io/gh/dgasmith/opt_einsum/pull/149?src=pr&el=h1) Report
> Merging [#149](https://codecov.io/gh/dgasmith/opt_einsum/pull/149?src=pr&el=desc) into [master](https://codecov.io/gh/dgasmith/opt_einsum/commit/2cbd28e612ff0af54656c34d7f6c239efe524b12&el=desc) will **decrease** coverage by `0.05%`.
> The diff coverage is `89.47%`.


jcmgray: OK, raised number of 'remaining terms' kept to 20 and should have fixed #147 now as well.

Patch:
diff --git a/opt_einsum/backends/jax.py b/opt_einsum/backends/jax.py
index 6bc1a1f..aefd941 100644
--- a/opt_einsum/backends/jax.py
+++ b/opt_einsum/backends/jax.py
@@ -13,18 +13,18 @@ _JAX = None
 
 
 def _get_jax_and_to_jax():
-  global _JAX
-  if _JAX is None:
-    import jax
+    global _JAX
+    if _JAX is None:
+        import jax
 
-    @to_backend_cache_wrap
-    @jax.jit
-    def to_jax(x):
-        return x
+        @to_backend_cache_wrap
+        @jax.jit
+        def to_jax(x):
+            return x
 
-    _JAX = jax, to_jax
+        _JAX = jax, to_jax
 
-  return _JAX
+    return _JAX
 
 
 def build_expression(_, expr):  # pragma: no cover
diff --git a/opt_einsum/contract.py b/opt_einsum/contract.py
index 6e7550a..1403f98 100644
--- a/opt_einsum/contract.py
+++ b/opt_einsum/contract.py
@@ -51,16 +51,22 @@ class PathInfo(object):
             "  Complete contraction:  {}\n".format(self.eq), "         Naive scaling:  {}\n".format(len(self.indices)),
             "     Optimized scaling:  {}\n".format(max(self.scale_list)), "      Naive FLOP count:  {:.3e}\n".format(
                 self.naive_cost), "  Optimized FLOP count:  {:.3e}\n".format(self.opt_cost),
-            "   Theoretical speedup:  {:3.3f}\n".format(self.speedup),
+            "   Theoretical speedup:  {:.3e}\n".format(self.speedup),
             "  Largest intermediate:  {:.3e} elements\n".format(self.largest_intermediate), "-" * 80 + "\n",
             "{:>6} {:>11} {:>22} {:>37}\n".format(*header), "-" * 80
         ]
 
         for n, contraction in enumerate(self.contraction_list):
             inds, idx_rm, einsum_str, remaining, do_blas = contraction
-            remaining_str = ",".join(remaining) + "->" + self.output_subscript
-            path_run = (self.scale_list[n], do_blas, einsum_str, remaining_str)
-            path_print.append("\n{:>4} {:>14} {:>22} {:>37}".format(*path_run))
+
+            if remaining is not None:
+                remaining_str = ",".join(remaining) + "->" + self.output_subscript
+            else:
+                remaining_str = "..."
+            size_remaining = max(0, 56 - max(22, len(einsum_str)))
+
+            path_run = (self.scale_list[n], do_blas, einsum_str, remaining_str, size_remaining)
+            path_print.append("\n{:>4} {:>14} {:>22}    {:>{}}".format(*path_run))
 
         return "".join(path_print)
 
@@ -303,7 +309,14 @@ def contract_path(*operands, **kwargs):
 
         einsum_str = ",".join(tmp_inputs) + "->" + idx_result
 
-        contraction = (contract_inds, idx_removed, einsum_str, input_list[:], do_blas)
+        # for large expressions saving the remaining terms at each step can
+        # incur a large memory footprint - and also be messy to print
+        if len(input_list) <= 20:
+            remaining = tuple(input_list)
+        else:
+            remaining = None
+
+        contraction = (contract_inds, idx_removed, einsum_str, remaining, do_blas)
         contraction_list.append(contraction)
 
     opt_cost = sum(cost_list)
@@ -529,7 +542,7 @@ def _core_contract(operands, contraction_list, backend='auto', evaluate_constant
 
     # Start contraction loop
     for num, contraction in enumerate(contraction_list):
-        inds, idx_rm, einsum_str, remaining, blas_flag = contraction
+        inds, idx_rm, einsum_str, _, blas_flag = contraction
 
         # check if we are performing the pre-pass of an expression with constants,
         #     if so, break out upon finding first non-constant (None) operand


Test patch:
diff --git a/opt_einsum/tests/test_contract.py b/opt_einsum/tests/test_contract.py
index e6dfb04..8268e88 100644
--- a/opt_einsum/tests/test_contract.py
+++ b/opt_einsum/tests/test_contract.py
@@ -171,7 +171,7 @@ def test_printing():
     views = helpers.build_views(string)
 
     ein = contract_path(string, *views)
-    assert len(str(ein[1])) == 726
+    assert len(str(ein[1])) == 728
 
 
 @pytest.mark.parametrize("string", tests)
diff --git a/opt_einsum/tests/test_input.py b/opt_einsum/tests/test_input.py
index e5c1677..f3b6c06 100644
--- a/opt_einsum/tests/test_input.py
+++ b/opt_einsum/tests/test_input.py
@@ -34,7 +34,8 @@ def test_type_errors():
         contract("", 0, out='test')
 
     # order parameter must be a valid order
-    with pytest.raises(TypeError):
+    # changed in Numpy 1.19, see https://github.com/numpy/numpy/commit/35b0a051c19265f5643f6011ee11e31d30c8bc4c
+    with pytest.raises((TypeError, ValueError)):
         contract("", 0, order='W')
 
     # casting parameter must be a valid casting