Skip to content
New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Add a LazyFunsor wrapper similar to LazyTensor #471

Open
fritzo opened this issue Mar 2, 2021 · 0 comments
Open

Add a LazyFunsor wrapper similar to LazyTensor #471

fritzo opened this issue Mar 2, 2021 · 0 comments
Labels
discussion enhancement New feature or request

Comments

@fritzo
Copy link
Member

fritzo commented Mar 2, 2021

This proposes a way to implement a LazyTensor-like feature in Funsor, should we decide we want it.

As explained in the LazyTensor paper, this interface improves debuggability by allowing code to run in different ways depending on whether or not attributes are accessed (e.g. via an assert statement or debug print statement), i.e. the interpretation is controlled not by an outer context manager but by inner attribute access. This might be a good interface for Funsor use in Pyro, where interpretations are outside of model-author control. E.g.

  x = pyro.sample("x", dist.Normal(0, 1))
+ print(x)  # <--- this could trigger eager evaluation
  pyro.sample("y", dist.Normal(x, 1), obs=data)

The idea of LazyFunsor would be to create a mutable object with a Funsor-like interface, to allow construction of reflected funsor terms, and then to trigger some other interpretation as soon as any attribute of that LazyFunsor is accessed, e.g. .data or maybe even .inputs. Here's a rough sketch

class LazyFunsor:
    def __init__(self, term: Funsor):
        self._term = term
        self._done = False

    # Any sort of access triggers interpretation.
    def __str__(self):
        return str(self._eval())

    def _eval(self):
        if not self._done:
            self._term = reinterpret(self._term)
            self._done = True
        return self._term

    # Operations are evaluated under reflect.
    def __add__(self, other):
        if isinstance(other, LazyFunsor):
            other = other._term
        with reflect:
            term = self._term + other
        return LazyFunsor(term, self._interpretation)

# We'd want to convert to_funsor before leaving the enclosing interpretation context.
@to_funsor.register(LazyFunsor)
def lazy_funsor_to_funsor(lazy_term):
    return lazy_term._eval()
@fritzo fritzo added enhancement New feature or request discussion labels Mar 2, 2021
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment
Labels
discussion enhancement New feature or request
Projects
None yet
Development

No branches or pull requests

1 participant