Add Precondition interpretation for Gaussian TVE

docs/source/interpretations.rst

@@ -23,6 +23,13 @@ Monte Carlo
     :show-inheritance:
     :member-order: bysource
 
+Preconditioning
+---------------
+.. automodule:: funsor.precondition
+    :members:
+    :show-inheritance:
+    :member-order: bysource
+
 Approximations
 --------------
 .. automodule:: funsor.approximations

funsor/__init__.py

@@ -46,6 +46,7 @@
     joint,
     montecarlo,
     ops,
+    precondition,
     recipes,
     sum_product,
     terms,
@@ -102,6 +103,7 @@
     "montecarlo",
     "of_shape",
     "ops",
+    "precondition",
     "pretty",
     "quote",
     "reals",

funsor/adjoint.py

@@ -67,10 +67,7 @@ def __enter__(self):
         self._old_interpretation = interpreter.get_interpretation()
         return super().__enter__()
 
-    def adjoint(self, sum_op, bin_op, root, targets=None, *, batch_vars=frozenset()):
-        # TODO Replace this with root + Constant(...) after #548 merges.
-        root_vars = root.input_vars | batch_vars
-
+    def adjoint(self, sum_op, bin_op, root, targets=None, *, batch_vars=set()):
         zero = to_funsor(ops.UNITS[sum_op])
         one = to_funsor(ops.UNITS[bin_op])
         adjoint_values = defaultdict(lambda: zero)
@@ -118,7 +115,7 @@ def adjoint(self, sum_op, bin_op, root, targets=None, *, batch_vars=frozenset())
             in_adjs = adjoint_ops(fn, sum_op, bin_op, adjoint_values[output], *inputs)
             for v, adjv in in_adjs:
                 # Marginalize out message variables that don't appear in recipients.
-                agg_vars = adjv.input_vars - v.input_vars - root_vars
+                agg_vars = adjv.input_vars - v.input_vars - root.input_vars - batch_vars
                 assert "particle" not in {var.name for var in agg_vars}  # DEBUG FIXME
                 old_value = adjoint_values[v]
                 adjoint_values[v] = sum_op(old_value, adjv.reduce(sum_op, agg_vars))

funsor/cnf.py

@@ -106,6 +106,11 @@ def _sample(self, sampled_vars, sample_inputs, rng_key):
         sampled_vars = sampled_vars.intersection(self.inputs)
         if not sampled_vars:
             return self
+        for term in self.terms:
+            if isinstance(term, Delta):
+                sampled_vars -= term.fresh
+        if not sampled_vars:
+            return self
 
         if self.red_op in (ops.null, ops.logaddexp):
             if rng_key is not None and get_backend() == "jax":
@@ -116,8 +121,8 @@ def _sample(self, sampled_vars, sample_inputs, rng_key):
                 rng_keys = [None] * len(self.terms)
 
             if self.bin_op in (ops.null, ops.logaddexp):
-                # Design choice: we sample over logaddexp reductions, but leave logaddexp
-                # binary choices symbolic.
+                # Design choice: we sample over logaddexp reductions, but leave
+                # logaddexp binary choices symbolic.
                 terms = [
                     term._sample(
                         sampled_vars.intersection(term.inputs), sample_inputs, rng_key
@@ -132,11 +137,15 @@ def _sample(self, sampled_vars, sample_inputs, rng_key):
                     greedy_vars = sampled_vars.intersection(term.inputs)
                     if greedy_vars:
                         break
+                assert greedy_vars
                 greedy_terms, terms = [], []
                 for term in self.terms:
-                    (
-                        terms if greedy_vars.isdisjoint(term.inputs) else greedy_terms
-                    ).append(term)
+                    if greedy_vars.isdisjoint(term.inputs):
+                        terms.append(term)
+                    elif isinstance(term, Delta) and greedy_vars.isdisjoint(term.fresh):
+                        terms.append(term)
+                    else:
+                        greedy_terms.append(term)
                 if len(greedy_terms) == 1:
                     term = greedy_terms[0]
                     terms.append(term._sample(greedy_vars, sample_inputs, rng_keys[0]))
@@ -392,7 +401,7 @@ def _(fn):
 # Normalizing Contractions
 ##########################################
 
-ORDERING = {Delta: 1, Number: 2, Tensor: 3, Gaussian: 4}
+ORDERING = {Delta: 1, Number: 2, Tensor: 3, Gaussian: 4, Unary[ops.NegOp, Gaussian]: 5}
 GROUND_TERMS = tuple(ORDERING)
 
 

funsor/domains.py

@@ -494,6 +494,22 @@ def _find_domain_stack(op, parts):
     return output
 
 
+@find_domain.register(ops.CatOp)
+def _find_domain_cat(op, parts):
+    dim = op.defaults["axis"]
+    if dim >= 0:
+        event_dims = {len(x.shape) for x in parts}
+        assert len(event_dims) == 1, "undefined"
+        dim = dim - next(iter(event_dims))
+    assert dim < 0
+    shape = broadcast_shape(*(x.shape[:dim] for x in parts))
+    shape += (sum(x.shape[dim] for x in parts),)
+    if dim < -1:
+        shape += broadcast_shape(*(x.shape[dim + 1 :] for x in parts))
+    output = Array[parts[0].dtype, shape]
+    return output
+
+
 @find_domain.register(ops.EinsumOp)
 def _find_domain_einsum(op, operands):
     equation = op.defaults["equation"]