Like memoized, but keep one cache per default graph.

def graph_memoized(func):
    """
    Like memoized, but keep one cache per default graph.
    """

    # TODO it keeps the graph alive
    from ..compat import tfv1
    GRAPH_ARG_NAME = '__IMPOSSIBLE_NAME_FOR_YOU__'

    @memoized
    def func_with_graph_arg(*args, **kwargs):
        kwargs.pop(GRAPH_ARG_NAME)
        return func(*args, **kwargs)

    @functools.wraps(func)
    def wrapper(*args, **kwargs):
        assert GRAPH_ARG_NAME not in kwargs, "No Way!!"
        graph = tfv1.get_default_graph()
        kwargs[GRAPH_ARG_NAME] = graph
        return func_with_graph_arg(*args, **kwargs)
    return wrapper