refactor memoize

Author: ferrine

pymc3/memoize.py

@@ -1,47 +1,59 @@
 import functools
 import pickle
-
+import collections
+from .util import biwrap
 CACHE_REGISTRY = []
 
 
-def memoize(obj):
+@biwrap
+def memoize(obj, bound=False):
     """
     An expensive memoizer that works with unhashables
     """
-    cache = obj.cache = {}
-    CACHE_REGISTRY.append(cache)
+    # this is declared not to be a bound method, so just attach new attr to obj
+    if not bound:
+        obj.cache = {}
+        CACHE_REGISTRY.append(obj.cache)
 
     @functools.wraps(obj)
     def memoizer(*args, **kwargs):
-        # remember first argument as well, used to clear cache for particular instance
-        key = (hashable(args[:1]), hashable(args), hashable(kwargs))
-
+        if not bound:
+            key = (hashable(args), hashable(kwargs))
+            cache = obj.cache
+        else:
+            # bound methods have self as first argument, remove it to compute key
+            key = (hashable(args[1:]), hashable(kwargs))
+            if not hasattr(args[0], '_cache'):
+                setattr(args[0], '_cache', collections.defaultdict(dict))
+                # do not add to cache regestry
+            cache = getattr(args[0], '_cache')[obj.__name__]
         if key not in cache:
             cache[key] = obj(*args, **kwargs)
 
         return cache[key]
     return memoizer
 
 
-def clear_cache():
-    for c in CACHE_REGISTRY:
-        c.clear()
+def clear_cache(obj=None):
+    if obj is None:
+        for c in CACHE_REGISTRY:
+            c.clear()
+    else:
+        if isinstance(obj, WithMemoization):
+            for v in getattr(obj, '_cache', {}).values():
+                v.clear()
+        else:
+            obj.cache.clear()
 
 
 class WithMemoization(object):
     def __hash__(self):
         return hash(id(self))
 
-    def __del__(self):
-        # regular property call with args (self, )
-        key = hash((self, ))
-        to_del = []
-        for c in CACHE_REGISTRY:
-            for k in c.keys():
-                if k[0] == key:
-                    to_del.append((c, k))
-        for (c, k) in to_del:
-            del c[k]
+    def __getstate__(self):
+        state = self.__dict__.copy()
+        state.pop('_cache', None)
+        return state
 
 
 def hashable(a):

pymc3/model.py

@@ -619,7 +619,7 @@ def isroot(self):
         return self.parent is None
 
     @property
-    @memoize
+    @memoize(bound=True)
     def bijection(self):
         vars = inputvars(self.cont_vars)
 

pymc3/util.py

@@ -1,5 +1,5 @@
 import re
-
+import functools
 from numpy import asscalar
 
 LATEX_ESCAPE_RE = re.compile(r'(%|_|\$|#|&)', re.MULTILINE)
@@ -152,3 +152,20 @@ def get_transformed(z):
     if hasattr(z, 'transformed'):
         z = z.transformed
     return z
+
+
+def biwrap(wrapper):
+    @functools.wraps(wrapper)
+    def enhanced(*args, **kwargs):
+        is_bound_method = hasattr(args[0], wrapper.__name__) if args else False
+        if is_bound_method:
+            count = 1
+        else:
+            count = 0
+        if len(args) > count:
+            newfn = wrapper(*args, **kwargs)
+            return newfn
+        else:
+            newwrapper = functools.partial(wrapper, *args, **kwargs)
+            return newwrapper
+    return enhanced

pymc3/variational/approximations.py

@@ -539,12 +539,6 @@ def __init__(self, *args, **kwargs):
     def __getattr__(self, item):
         return getattr(self.groups[0], item)
 
-    def __getstate__(self):
-        return self.__dict__.copy()
-
-    def __setstate__(self, state):
-        self.__dict__.update(state)
-
 
 class MeanField(SingleGroupApproximation):
     __doc__ = """**Single Group Mean Field Approximation**

pymc3/variational/opvi.py

@@ -106,10 +106,10 @@ def node_property(f):
     if isinstance(f, str):
 
         def wrapper(fn):
-            return property(memoize(change_flags(compute_test_value='off')(append_name(f)(fn))))
+            return property(memoize(change_flags(compute_test_value='off')(append_name(f)(fn)), bound=True))
         return wrapper
     else:
-        return property(memoize(change_flags(compute_test_value='off')(f)))
+        return property(memoize(change_flags(compute_test_value='off')(f), bound=True))
 
 
 @change_flags(compute_test_value='ignore')
@@ -1487,7 +1487,7 @@ def vars_names(vs):
         return found
 
     @property
-    @memoize
+    @memoize(bound=True)
     @change_flags(compute_test_value='off')
     def sample_dict_fn(self):
         s = tt.iscalar()