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
```py
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()
```
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.
The idea of
LazyFunsorwould 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 thatLazyFunsoris accessed, e.g..dataor maybe even.inputs. Here's a rough sketch