ProvenanceTensor

funsor/adjoint.py

@@ -58,7 +58,8 @@ def interpret(self, cls, *args):
             )
             for arg in args
         ]
-        self._eager_to_lazy[result] = reflect.interpret(cls, *lazy_args)
+        #  with self._old_interpretation:
+        #      self._eager_to_lazy[result] = reflect.interpret(cls, *lazy_args)
         return result
 
     def __enter__(self):
@@ -83,34 +84,34 @@ def adjoint(self, sum_op, bin_op, root, targets=None):
                     continue
 
             # reverse the effects of alpha-renaming
-            with reflect:
-
-                lazy_output = self._eager_to_lazy[output]
-                lazy_fn = type(lazy_output)
-                lazy_inputs = lazy_output._ast_values
-                # TODO abstract this into a helper function
-                # FIXME make lazy_output linear instead of quadratic in the size of the tape
-                lazy_other_subs = tuple(
-                    (name, to_funsor(name.split("__BOUND")[0], domain))
-                    for name, domain in lazy_output.inputs.items()
-                    if "__BOUND" in name
-                )
-                lazy_inputs = _alpha_unmangle(
-                    substitute(lazy_fn(*lazy_inputs), lazy_other_subs)
-                )
-                lazy_output = type(lazy_output)(
-                    *_alpha_unmangle(substitute(lazy_output, lazy_other_subs))
-                )
-
-                other_subs = tuple(
-                    (name, to_funsor(name.split("__BOUND")[0], domain))
-                    for name, domain in output.inputs.items()
-                    if "__BOUND" in name
-                )
-                inputs = _alpha_unmangle(substitute(fn(*inputs), other_subs))
-                output = type(output)(*_alpha_unmangle(substitute(output, other_subs)))
-
-                self._eager_to_lazy[output] = lazy_output
+            #  with reflect:
+            #
+            #      lazy_output = self._eager_to_lazy[output]
+            #      lazy_fn = type(lazy_output)
+            #      lazy_inputs = lazy_output._ast_values
+            #      # TODO abstract this into a helper function
+            #      # FIXME make lazy_output linear instead of quadratic in the size of the tape
+            #      lazy_other_subs = tuple(
+            #          (name, to_funsor(name.split("__BOUND")[0], domain))
+            #          for name, domain in lazy_output.inputs.items()
+            #          if "__BOUND" in name
+            #      )
+            #      lazy_inputs = _alpha_unmangle(
+            #          substitute(lazy_fn(*lazy_inputs), lazy_other_subs)
+            #      )
+            #      lazy_output = type(lazy_output)(
+            #          *_alpha_unmangle(substitute(lazy_output, lazy_other_subs))
+            #      )
+            #
+            #      other_subs = tuple(
+            #          (name, to_funsor(name.split("__BOUND")[0], domain))
+            #          for name, domain in output.inputs.items()
+            #          if "__BOUND" in name
+            #      )
+            #      inputs = _alpha_unmangle(substitute(fn(*inputs), other_subs))
+            #      output = type(output)(*_alpha_unmangle(substitute(output, other_subs)))
+            #
+            #      self._eager_to_lazy[output] = lazy_output
 
             in_adjs = adjoint_ops(fn, sum_op, bin_op, adjoint_values[output], *inputs)
             for v, adjv in in_adjs:

funsor/cnf.py

@@ -102,7 +102,9 @@ def __str__(self):
             )
         return super().__str__()
 
-    def unscaled_sample(self, sampled_vars, sample_inputs, rng_key=None):
+    def unscaled_sample(
+        self, sampled_vars, sample_inputs, rng_key=None, raw_value=None
+    ):
         sampled_vars = sampled_vars.intersection(self.inputs)
         if not sampled_vars:
             return self
@@ -120,7 +122,9 @@ def unscaled_sample(self, sampled_vars, sample_inputs, rng_key=None):
                 # binary choices symbolic.
                 terms = [
                     term.unscaled_sample(
-                        sampled_vars.intersection(term.inputs), sample_inputs
+                        sampled_vars.intersection(term.inputs),
+                        sample_inputs,
+                        raw_value=raw_value,
                     )
                     for term, rng_key in zip(self.terms, rng_keys)
                 ]
@@ -141,13 +145,19 @@ def unscaled_sample(self, sampled_vars, sample_inputs, rng_key=None):
                         break
                 greedy_terms, terms = [], []
                 for term in self.terms:
+                    if isinstance(term, funsor.distribution.Distribution):
+                        term_var = {term.value.name}
+                    else:
+                        term_var = term.inputs
                     (
-                        terms if greedy_vars.isdisjoint(term.inputs) else greedy_terms
+                        terms if greedy_vars.isdisjoint(term_var) else greedy_terms
                     ).append(term)
                 if len(greedy_terms) == 1:
                     term = greedy_terms[0]
                     terms.append(
-                        term.unscaled_sample(greedy_vars, sample_inputs, rng_keys[0])
+                        term.unscaled_sample(
+                            greedy_vars, sample_inputs, rng_keys[0], raw_value=raw_value
+                        )
                     )
                     result = Contraction(
                         self.red_op, self.bin_op, self.reduced_vars, *terms
@@ -162,7 +172,9 @@ def unscaled_sample(self, sampled_vars, sample_inputs, rng_key=None):
                     terms.append(gaussian)
                     terms.append(-gaussian.log_normalizer)
                     terms.append(
-                        term.unscaled_sample(greedy_vars, sample_inputs, rng_keys[0])
+                        term.unscaled_sample(
+                            greedy_vars, sample_inputs, rng_keys[0], raw_value=raw_value
+                        )
                     )
                     result = Contraction(
                         self.red_op, self.bin_op, self.reduced_vars, *terms
@@ -174,7 +186,9 @@ def unscaled_sample(self, sampled_vars, sample_inputs, rng_key=None):
                 ):
                     sampled_terms = [
                         term.unscaled_sample(
-                            greedy_vars.intersection(term.value.inputs), sample_inputs
+                            greedy_vars.intersection(term.value.inputs),
+                            sample_inputs,
+                            raw_value=raw_value,
                         )
                         for term in greedy_terms
                         if isinstance(term, funsor.distribution.Distribution)
@@ -193,7 +207,10 @@ def unscaled_sample(self, sampled_vars, sample_inputs, rng_key=None):
                         )
                     )
                 return result.unscaled_sample(
-                    sampled_vars - greedy_vars, sample_inputs, rng_keys[1]
+                    sampled_vars - greedy_vars,
+                    sample_inputs,
+                    rng_keys[1],
+                    raw_value=raw_value,
                 )
 
         raise TypeError(

funsor/constant.py

@@ -0,0 +1,128 @@
+# Copyright Contributors to the Pyro project.
+# SPDX-License-Identifier: Apache-2.0
+
+from collections import OrderedDict
+
+from funsor.distribution import Distribution
+from funsor.tensor import Tensor
+from funsor.terms import (
+    Binary,
+    Funsor,
+    FunsorMeta,
+    Number,
+    Unary,
+    Variable,
+    eager,
+    to_data,
+    to_funsor,
+)
+from funsor.torch.provenance import ProvenanceTensor
+
+from .ops import BinaryOp, UnaryOp
+
+
+class ConstantMeta(FunsorMeta):
+    """
+    Wrapper to convert ``const_inputs`` to a tuple.
+    """
+
+    def __call__(cls, const_inputs, arg):
+        if isinstance(const_inputs, dict):
+            const_inputs = tuple(const_inputs.items())
+
+        return super(ConstantMeta, cls).__call__(const_inputs, arg)
+
+
+class Constant(Funsor, metaclass=ConstantMeta):
+    def __init__(self, const_inputs, arg):
+        assert isinstance(arg, Funsor)
+        assert isinstance(const_inputs, tuple)
+        assert set(const_inputs).isdisjoint(arg.inputs)
+        const_inputs = OrderedDict(const_inputs)
+        inputs = const_inputs.copy()
+        inputs.update(arg.inputs)
+        output = arg.output
+        fresh = frozenset(const_inputs)
+        bound = {}
+        super(Constant, self).__init__(inputs, output, fresh, bound)
+        self.arg = arg
+        self.const_vars = frozenset(Variable(k, v) for k, v in const_inputs.items())
+        self.const_inputs = const_inputs
+
+    def eager_subs(self, subs):
+        assert isinstance(subs, tuple)
+        subs = OrderedDict((k, v) for k, v in subs)
+        const_inputs = OrderedDict()
+        for k, d in self.const_inputs.items():
+            # handle when subs is in self.arg.inputs
+            if k in subs:
+                v = subs[k]
+                if isinstance(v, Variable):
+                    del subs[k]
+                    k = v.name
+                    const_inputs[k] = d
+        if const_inputs:
+            return Constant(const_inputs, self.arg)
+        return self.arg
+
+    def eager_reduce(self, op, reduced_vars):
+        assert reduced_vars.issubset(self.inputs)
+        const_inputs = OrderedDict(
+            (k, v) for k, v in self.const_inputs.items() if k not in reduced_vars
+        )
+        reduced_vars = reduced_vars - frozenset(self.const_inputs)
+        reduced_arg = self.arg.reduce(op, reduced_vars)
+        if const_inputs:
+            return Constant(const_inputs, reduced_arg)
+        return reduced_arg
+
+
+@eager.register(Binary, BinaryOp, Constant, Constant)
+def eager_binary_constant_constant(op, lhs, rhs):
+    const_inputs = OrderedDict(
+        (k, v) for k, v in lhs.const_inputs.items() if k not in rhs.const_inputs
+    )
+    const_inputs.update(
+        (k, v) for k, v in rhs.const_inputs.items() if k not in lhs.const_inputs
+    )
+    if const_inputs:
+        return Constant(const_inputs, op(lhs.arg, rhs.arg))
+    return op(lhs.arg, rhs.arg)
+
+
+@eager.register(Binary, BinaryOp, Constant, (Number, Tensor, Distribution))
+def eager_binary_constant_tensor(op, lhs, rhs):
+    const_inputs = OrderedDict(
+        (k, v) for k, v in lhs.const_inputs.items() if k not in rhs.inputs
+    )
+    if const_inputs:
+        return Constant(const_inputs, op(lhs.arg, rhs))
+    return op(lhs.arg, rhs)
+
+
+@eager.register(Binary, BinaryOp, (Number, Tensor, Distribution), Constant)
+def eager_binary_tensor_constant(op, lhs, rhs):
+    const_inputs = OrderedDict(
+        (k, v) for k, v in rhs.const_inputs.items() if k not in lhs.inputs
+    )
+    if const_inputs:
+        return Constant(const_inputs, op(lhs, rhs.arg))
+    return op(lhs, rhs.arg)
+
+
+@eager.register(Unary, UnaryOp, Constant)
+def eager_binary_tensor_constant(op, arg):
+    return Constant(arg.const_inputs, op(arg.arg))
+
+
+@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_funsor.register(ProvenanceTensor)
+def provenance_to_funsor(x, output=None, dim_to_name=None):
+    if isinstance(x, ProvenanceTensor):
+        ret = to_funsor(x._t, output=output, dim_to_name=dim_to_name)
+        return Constant(OrderedDict(x._provenance), ret)

funsor/delta.py

@@ -200,7 +200,9 @@ def eager_reduce(self, op, reduced_vars):
 
         return None  # defer to default implementation
 
-    def unscaled_sample(self, sampled_vars, sample_inputs, rng_key=None):
+    def unscaled_sample(
+        self, sampled_vars, sample_inputs, rng_key=None, raw_value=None
+    ):
         return self
 
 

funsor/distribution.py

@@ -12,6 +12,7 @@
 
 import makefun
 
+import funsor
 import funsor.delta
 import funsor.ops as ops
 from funsor.affine import is_affine
@@ -151,6 +152,11 @@ def eager_reduce(self, op, reduced_vars):
             and isinstance(self.value, Variable)
             and self.value.name in reduced_vars
         ):
+            const_inputs = OrderedDict(
+                (k, v) for k, v in self.inputs.items() if k not in reduced_vars
+            )
+            if const_inputs:
+                return funsor.constant.Constant(const_inputs, Number(0.0))
             return Number(0.0)  # distributions are normalized
         return super(Distribution, self).eager_reduce(op, reduced_vars)
 
@@ -206,7 +212,9 @@ def eager_log_prob(cls, *params):
             inputs.update(x.inputs)
         return log_prob.align(tuple(inputs))
 
-    def unscaled_sample(self, sampled_vars, sample_inputs, rng_key=None):
+    def unscaled_sample(
+        self, sampled_vars, sample_inputs, rng_key=None, raw_value=None
+    ):
 
         # note this should handle transforms correctly via distribution_to_data
         raw_dist, value_name, value_output, dim_to_name = self._get_raw_dist()
@@ -220,10 +228,14 @@ def unscaled_sample(self, sampled_vars, sample_inputs, rng_key=None):
         sample_args = (
             (sample_shape,) if get_backend() == "torch" else (rng_key, sample_shape)
         )
-        if raw_dist.has_rsample:
-            raw_value = raw_dist.rsample(*sample_args)
+
+        if raw_value is not None and value_name in raw_value:
+            raw_value = raw_value[value_name]
         else:
-            raw_value = ops.detach(raw_dist.sample(*sample_args))
+            if raw_dist.has_rsample:
+                raw_value = raw_dist.rsample(*sample_args)
+            else:
+                raw_value = ops.detach(raw_dist.sample(*sample_args))
 
         funsor_value = to_funsor(
             raw_value, output=value_output, dim_to_name=dim_to_name
@@ -232,6 +244,7 @@ def unscaled_sample(self, sampled_vars, sample_inputs, rng_key=None):
             tuple(sample_inputs)
             + tuple(inp for inp in self.inputs if inp in funsor_value.inputs)
         )
+
         result = funsor.delta.Delta(value_name, funsor_value)
         if not raw_dist.has_rsample:
             # scaling of dice_factor by num samples should already be handled by Funsor.sample

funsor/montecarlo.py

@@ -19,9 +19,10 @@ class MonteCarlo(StatefulInterpretation):
     :param rng_key:
     """
 
-    def __init__(self, *, rng_key=None, **sample_inputs):
+    def __init__(self, *, rng_key=None, raw_value=None, **sample_inputs):
         super().__init__("monte_carlo")
         self.rng_key = rng_key
+        self.raw_value = raw_value
         self.sample_inputs = OrderedDict(sample_inputs)
 
 
@@ -33,7 +34,9 @@ def monte_carlo_integrate(state, log_measure, integrand, reduced_vars):
 
         sample_options["rng_key"], state.rng_key = jax.random.split(state.rng_key)
 
-    sample = log_measure.sample(reduced_vars, state.sample_inputs, **sample_options)
+    sample = log_measure.sample(
+        reduced_vars, state.sample_inputs, raw_value=state.raw_value, **sample_options
+    )
     if sample is log_measure:
         return None  # cannot progress
     reduced_vars |= frozenset(

funsor/tensor.py

@@ -329,7 +329,9 @@ def eager_reduce(self, op, reduced_vars):
             return Tensor(data, inputs, dtype)
         return super(Tensor, self).eager_reduce(op, reduced_vars)
 
-    def unscaled_sample(self, sampled_vars, sample_inputs, rng_key=None):
+    def unscaled_sample(
+        self, sampled_vars, sample_inputs, rng_key=None, raw_value=None
+    ):
         assert self.output == Real
         sampled_vars = sampled_vars.intersection(self.inputs)
         if not sampled_vars: