Provenance funsor

docs/source/funsors.rst

@@ -64,3 +64,11 @@ Constant
     :undoc-members:
     :show-inheritance:
     :member-order: bysource
+
+Provenance
+----------
+.. automodule:: funsor.provenance
+    :members:
+    :undoc-members:
+    :show-inheritance:
+    :member-order: bysource

funsor/__init__.py

@@ -7,6 +7,7 @@
 from funsor.integrate import Integrate
 from funsor.interpreter import interpretation, reinterpret
 from funsor.op_factory import make_op
+from funsor.provenance import Provenance
 from funsor.sum_product import MarkovProduct
 from funsor.tensor import Tensor, function
 from funsor.terms import (
@@ -47,6 +48,7 @@
     montecarlo,
     ops,
     precondition,
+    provenance,
     recipes,
     sum_product,
     terms,
@@ -71,6 +73,7 @@
     "Number",
     "Real",
     "Reals",
+    "Provenance",
     "Slice",
     "Stack",
     "Tensor",
@@ -105,6 +108,7 @@
     "ops",
     "precondition",
     "pretty",
+    "provenance",
     "quote",
     "reals",
     "recipes",

funsor/distribution.py

@@ -19,6 +19,7 @@
 from funsor.domains import Array, Real, Reals
 from funsor.gaussian import Gaussian
 from funsor.interpreter import gensym
+from funsor.provenance import Provenance
 from funsor.tensor import (
     Tensor,
     align_tensors,
@@ -458,6 +459,14 @@ def backenddist_to_funsor(
         for param_name in funsor_dist_class._ast_fields
         if param_name != "value"
     ]
+    provenance = frozenset().union(
+        *[param.provenance for param in params if isinstance(param, Provenance)]
+    )
+    if provenance:
+        params = [
+            param.term if isinstance(param, Provenance) else param for param in params
+        ]
+        return Provenance(funsor_dist_class(*params), provenance)
     return funsor_dist_class(*params)
 
 

funsor/integrate.py

@@ -11,6 +11,7 @@
 from funsor.delta import Delta
 from funsor.gaussian import Gaussian, _norm2, _vm, align_gaussian
 from funsor.interpretations import eager, normalize
+from funsor.provenance import Provenance
 from funsor.tensor import Tensor
 from funsor.terms import (
     Funsor,
@@ -139,6 +140,7 @@ def normalize_integrate_contraction(log_measure, integrand, reduced_vars):
         Tensor,
         GaussianMixture,
         EagerConstant,
+        Provenance,
     ),
 )
 def eager_contraction_binary_to_integrate(red_op, bin_op, reduced_vars, lhs, rhs):

funsor/provenance.py

@@ -0,0 +1,105 @@
+# Copyright Contributors to the Pyro project.
+# SPDX-License-Identifier: Apache-2.0
+
+from collections import OrderedDict
+
+import funsor.ops as ops
+from funsor.tensor import Tensor
+from funsor.terms import Binary, Funsor, FunsorMeta, Number, Unary, Variable, eager
+
+
+class ProvenanceMeta(FunsorMeta):
+    """
+    Wrapper to combine provenance information from the term.
+    """
+
+    def __call__(cls, term, provenance):
+        while isinstance(term, Provenance):
+            provenance |= term.provenance
+            term = term.term
+
+        return super(ProvenanceMeta, cls).__call__(term, provenance)
+
+
+class Provenance(Funsor, metaclass=ProvenanceMeta):
+    """
+    Provenance funsor for tracking the dependence of terms on ``(name, point)``
+    of sampled random variables.
+
+    **References**
+
+    [1] David Wingate, Noah Goodman, Andreas Stuhlmüller, Jeffrey Siskind (2011)
+        Nonstandard Interpretations of Probabilistic Programs for Efficient Inference
+        http://papers.neurips.cc/paper/4309-nonstandard-interpretations-of-probabilistic-programs-for-efficient-inference.pdf
+
+    :param funsor term: A term that depends on tracked variables.
+    :param frozenset provenance: A set of tuples of the form ``(name, point)``.
+    """
+
+    def __init__(self, term, provenance):
+        assert isinstance(term, Funsor)
+        assert isinstance(provenance, frozenset)
+
+        provenance_names = frozenset([name for name, point in provenance])
+        assert provenance_names.isdisjoint(term.inputs)
+        inputs = OrderedDict()
+        for name, point in provenance:
+            assert isinstance(name, str)
+            assert isinstance(point, Funsor)
+            assert name not in point.inputs
+            inputs.update({name: point.output})
+            inputs.update(point.inputs)
+
+        inputs.update(term.inputs)
+        output = term.output
+        fresh = provenance_names
+        bound = {}
+        super(Provenance, self).__init__(inputs, output, fresh, bound)
+        self.term = term
+        self.provenance = provenance
+
+    def eager_subs(self, subs):
+        assert isinstance(subs, tuple)
+        subs = OrderedDict(subs)
+        assert set(subs).issubset(self.fresh)
+        new_provenance = frozenset()
+        new_term = self.term
+        for name, point in self.provenance:
+            if name in subs:
+                value = subs[name]
+                if isinstance(value, Variable):
+                    new_provenance |= frozenset([(value.name, point)])
+                    continue
+
+                # leave out the substituted provenance variable
+                # make sure that the value matches the point
+                assert value is point
+            else:
+                new_provenance |= frozenset([(name, point)])
+        return Provenance(new_term, new_provenance) if new_provenance else new_term
+
+    def _sample(self, sampled_vars, sample_inputs, rng_key):
+        result = self.term._sample(sampled_vars, sample_inputs, rng_key)
+        return Provenance(result, self.provenance)
+
+
+@eager.register(Binary, ops.BinaryOp, Provenance, Provenance)
+def eager_binary_provenance_provenance(op, lhs, rhs):
+    return Provenance(op(lhs.term, rhs.term), lhs.provenance | rhs.provenance)
+
+
+@eager.register(Binary, ops.BinaryOp, Provenance, (Number, Tensor))
+def eager_binary_provenance_tensor(op, lhs, rhs):
+    assert lhs.fresh.isdisjoint(rhs.inputs)
+    return Provenance(op(lhs.term, rhs), lhs.provenance)
+
+
+@eager.register(Binary, ops.BinaryOp, (Number, Tensor), Provenance)
+def eager_binary_tensor_provenance(op, lhs, rhs):
+    assert rhs.fresh.isdisjoint(lhs.inputs)
+    return Provenance(op(lhs, rhs.term), rhs.provenance)
+
+
+@eager.register(Unary, ops.UnaryOp, Provenance)
+def eager_unary(op, arg):
+    return Provenance(op(arg.term), arg.provenance)

funsor/torch/__init__.py

@@ -1,12 +1,10 @@
 # Copyright Contributors to the Pyro project.
 # SPDX-License-Identifier: Apache-2.0
 
-from collections import OrderedDict
-
 import torch
 from multipledispatch import dispatch
 
-from funsor.constant import Constant
+from funsor.provenance import Provenance
 from funsor.tensor import tensor_to_funsor
 from funsor.terms import to_data, to_funsor
 from funsor.torch.provenance import ProvenanceTensor
@@ -31,15 +29,15 @@ def _quote(x, indent, out):
 
 
 @to_funsor.register(ProvenanceTensor)
-def provenance_to_funsor(x, output=None, dim_to_name=None):
-    ret = to_funsor(x._t, output=output, dim_to_name=dim_to_name)
-    return Constant(OrderedDict(x._provenance), ret)
+def provenancetensor_to_funsor(x, output=None, dim_to_name=None):
+    term = to_funsor(x._t, output=output, dim_to_name=dim_to_name)
+    return Provenance(term, x._provenance)
 
 
-@to_data.register(Constant)
-def constant_to_data(x, name_to_dim=None):
-    data = to_data(x.arg, name_to_dim=name_to_dim)
-    return ProvenanceTensor(data, provenance=frozenset(x.const_inputs.items()))
+@to_data.register(Provenance)
+def provenance_to_data(x, name_to_dim=None):
+    data = to_data(x.term, name_to_dim=name_to_dim)
+    return ProvenanceTensor(data, provenance=x.provenance)
 
 
 to_funsor.register(torch.Tensor)(tensor_to_funsor)

test/test_constant.py

@@ -8,8 +8,8 @@
 from funsor.delta import Delta
 from funsor.domains import Bint, Real
 from funsor.tensor import Tensor
-from funsor.terms import Number, Variable, to_data, to_funsor
-from funsor.testing import assert_close, randn, requires_backend
+from funsor.terms import Number, Variable
+from funsor.testing import assert_close, randn
 
 
 def test_eager_subs_variable():
@@ -81,29 +81,3 @@ def test_align():
         for i in range(2):
             for j in range(3):
                 assert x(a=0, b=b, i=i, j=j) == y(a=0, b=b, i=i, j=j)
-
-
-@requires_backend("torch", reason="requires ProvenanceTensor")
-def test_to_funsor():
-    import torch
-
-    from funsor.torch.provenance import ProvenanceTensor
-
-    data = torch.zeros(3, 3)
-    pt = ProvenanceTensor(data, frozenset({("x", Real)}))
-    c = to_funsor(pt)
-    assert c is Constant(OrderedDict(x=Real), Tensor(data))
-
-
-@requires_backend("torch", reason="requires ProvenanceTensor")
-def test_to_data():
-    import torch
-
-    from funsor.torch.provenance import ProvenanceTensor
-
-    data = torch.zeros(3, 3)
-    c = Constant(OrderedDict(x=Real), Tensor(data))
-    pt = to_data(c)
-    assert isinstance(pt, ProvenanceTensor)
-    assert pt._t is data
-    assert pt._provenance == frozenset({("x", Real)})