Adds missing caching

tensorly · Nov 14, 2022 · 8000f98 · 8000f98
1 parent 993ae4f
commit 8000f98
Showing 1 changed file with 20 additions and 0 deletions.
diff --git a/tensorly/tenalg/einsum_tenalg/caching.py b/tensorly/tenalg/einsum_tenalg/caching.py
@@ -0,0 +1,20 @@
+EINSUM_PATH_CACHE = dict()
+
+
+def einsum_path_cached(fun):
+    def wrapped(key, *args, **kwargs):
+        name = fun.__name__
+        try:
+            cache = EINSUM_PATH_CACHE[name]
+        except KeyError:
+            EINSUM_PATH_CACHE[name] = dict()
+            cache = EINSUM_PATH_CACHE[name]
+        try:
+            equation = cache[key]
+        except KeyError:
+            equation = fun(*args, **kwargs)
+            cache[key] = equation
+
+        return equation
+
+    return wrapped