Python - how do I memoize a partial object?

Question

I have a set of functions that take integers and functions as arguments. I'd like to memoize them.

I know that using this solution, I could use pickle to encode both sets of arguments and memoize the encoded values. In this particular use case, however, the function arguments are large and constant, and I'd rather not take up the lru_cache space with a function argument that, in the course of a program run, never changes.

Is there a way for me to memoize a partial function, where I've fixed the function arguments and have received a partial object that takes only hashable arguments? I can't figure out how to use the functools.lru_cache decorator as a function.

Here's what I've tried on a toy example. It doesn't work; the binomial tree still revisits nodes.

import functools
import logging


logging.basicConfig(level=logging.DEBUG)


def binomial_tree(x, y, fn):
    logging.debug(f"binomial_tree({x}, {y})")
    """Note: this does not recombine, and we can't memoize function."""
    if x == 10:
        return fn(x, y)
    else:
        return 0.5 * binomial_tree(x + 1, y, fn) + 0.5 * binomial_tree(x + 1, y + 1, fn)


def memoize_fn(fn):
    @functools.lru_cache(maxsize=None)
    def inner(*args, **kwargs):
        return fn(*args, **kwargs)
    return inner

memoized_binomial_tree = memoize_fn(functools.partial(binomial_tree, fn=lambda x, y: 10 * x * y))
print(memoized_binomial_tree(0, 0))

The problem here is that your memoization of partial object doesn't affect the two calls `binomial_tree(x + 1, y, fn)` and `binomial_tree(x + 1, y + 1, fn)`. — sanyassh, Sep 14 '19 at 12:42
Right. I think I I'd be fine if I just memoized `binomial_tree`. But I don't want to in this instance. — MikeRand, Sep 14 '19 at 12:49
The function argument really never changes during a program run? It may be a stupid question but what's the point of making a partial instead of using a class or a global variable for the function? Or am I taking the toy example too serious? — MSeifert, Sep 14 '19 at 13:02
Library where I want to pass to select amongst multiple different versions of `fn`. But yes, changing a global would work. I just dont like globals. — MikeRand, Sep 14 '19 at 13:04

score 4 · Accepted Answer · answered Sep 14 '19 at 12:53

Here is a way to memoize your toy example with binomial_tree without encoding and memoizing function arguments:

import functools
import logging


logging.basicConfig(level=logging.DEBUG)


def create_binomial_tree(fn):
    @functools.lru_cache(maxsize=None)
    def binomial_tree(x, y):
        logging.debug(f"binomial_tree({x}, {y})")
        if x == 10:
            return fn(x, y)
        else:
            return 0.5 * binomial_tree(x + 1, y) + 0.5 * binomial_tree(x + 1, y + 1)
    return binomial_tree


memoized_binomial_tree = create_binomial_tree(fn=lambda x, y: 10 * x * y)
print(memoized_binomial_tree(0, 0))

Maybe it can be applicable in your real use case?

Ok, so I mimic partial with closure but insert the caching in between. That works. — MikeRand, Sep 14 '19 at 12:55

Python - how do I memoize a partial object?

1 Answers1