Disable torch.backends.opt_einsum to avoid duplicate work

opt_einsum/backends/torch.py

@@ -6,6 +6,7 @@
 
 from ..parser import convert_to_valid_einsum_chars
 from ..sharing import to_backend_cache_wrap
+from functools import lru_cache
 
 __all__ = [
     "transpose",
@@ -41,14 +42,29 @@ def transpose(a, axes):
     return a.permute(*axes)
 
 
+@lru_cache(None)
+def get_einsum():
+    torch, _ = _get_torch_and_device()
+    if hasattr(torch, "_VF") and hasattr(torch._VF, "einsum"):
+        print("called")
+        return torch._VF.einsum
+    print("did not get called")
+
+    def einsum_no_opt(*args, **kwargs):
+        if hasattr(torch, "backends") and hasattr(torch.backends, "opt_einsum"):
+            with torch.backends.opt_einsum.flags(enabled=False):
+                return torch.einsum(*args, **kwargs)
+        return torch.einsum(*args, **kwargs)
+
+    return einsum_no_opt
+
+
 def einsum(equation, *operands):
     """Variadic version of torch.einsum to match numpy api."""
     # rename symbols to support PyTorch 0.4.1 and earlier,
     # which allow only symbols a-z.
     equation = convert_to_valid_einsum_chars(equation)
-
-    torch, _ = _get_torch_and_device()
-    return torch.einsum(equation, operands)
+    return get_einsum()(equation, operands)
 
 
 def tensordot(x, y, axes=2):

opt_einsum/tests/test_contract.py

@@ -1,5 +1,5 @@
 """
-Tets a series of opt_einsum contraction paths to ensure the results are the same for different paths
+Tests a series of opt_einsum contraction paths to ensure the results are the same for different paths
 """
 
 import numpy as np
@@ -198,6 +198,17 @@ def test_contract_expression_interleaved_input():
     assert np.allclose(out, expected)
 
 
+def test_torch_contract_expression_interleaved_input():
+    import torch
+
+    x, y, z = (torch.randn(2, 2) for _ in "xyz")
+    expected = np.einsum(x, [0, 1], y, [1, 2], z, [2, 3], [3, 0])
+    xshp, yshp, zshp = ((2, 2) for _ in "xyz")
+    expr = contract_expression(xshp, [0, 1], yshp, [1, 2], zshp, [2, 3], [3, 0])
+    out = expr(x, y, z)
+    assert np.allclose(out, expected)
+
+
 @pytest.mark.parametrize(
     "string,constants",
     [