Cache on Operator

Author: jrmaddison

pyadjoint/adjfloat.py

@@ -1,4 +1,4 @@
-from functools import cached_property, lru_cache, wraps
+from functools import cached_property, wraps
 from itertools import count
 import numbers
 import operator
@@ -45,7 +45,6 @@ def fdiff(self, argindex=1):
             return sp.Integer(1) / (sp.Integer(1) + self.args[0])
 
 
-@lru_cache(maxsize=128)
 def codegen(expr, symbols, diff=()):
     for idx in diff:
         expr = expr.diff(symbols[idx])
@@ -58,6 +57,7 @@ class Operator:
     def __init__(self, sp_operator, nargs):
         self._sp_operator = sp_operator
         self._nargs = nargs
+        self._cache = {}
 
     @property
     def sp_operator(self):
@@ -76,7 +76,9 @@ def expr(self):
         return self.sp_operator(*self.symbols)
 
     def codegen(self, diff=()):
-        return codegen(self.expr, self.symbols, diff=diff)
+        if diff not in self._cache:
+            self._cache[diff] = codegen(self.expr, self.symbols, diff=diff)
+        return self._cache[diff]
 
 
 class AdjFloatExprBlock(Block):