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

@@ -44,6 +44,7 @@
     joint,
     montecarlo,
     ops,
+    precondition,
     recipes,
     sum_product,
     terms,
@@ -98,6 +99,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/affine.py

@@ -61,7 +61,9 @@ def _(fn):
 
 @affine_inputs.register(Unary)
 def _(fn):
-    if fn.op in (ops.neg, ops.sum) or isinstance(fn.op, ops.ReshapeOp):
+    if fn.op in (ops.neg, ops.sum) or isinstance(
+        fn.op, (ops.ReshapeOp, ops.GetsliceOp)
+    ):
         return affine_inputs(fn.arg)
     return frozenset()
 

funsor/domains.py

@@ -10,6 +10,7 @@
 from weakref import WeakValueDictionary
 
 import funsor.ops as ops
+from funsor.ops.builtin import parse_ellipsis, parse_slice
 from funsor.util import broadcast_shape, get_backend, get_tracing_state, quote
 
 Domain = type
@@ -331,6 +332,58 @@ def _find_domain_getitem(op, lhs_domain, rhs_domain):
         )
 
 
+@find_domain.register(ops.GetsliceOp)
+def _find_domain_getslice(op, domain):
+    index = op.defaults["index"]
+    if isinstance(domain, ArrayType):
+        dtype = domain.dtype
+        shape = list(domain.shape)
+        left, right = parse_ellipsis(index)
+
+        i = 0
+        for part in left:
+            if part is None:
+                shape.insert(i, 1)
+                i += 1
+            elif isinstance(part, int):
+                del shape[i]
+            elif isinstance(part, slice):
+                start, stop, step = parse_slice(part, shape[i])
+                shape[i] = max(0, (stop - start + step - 1) // step)
+                i += 1
+            else:
+                raise ValueError(part)
+
+        i = -1
+        for part in reversed(right):
+            if part is None:
+                shape.insert(len(shape) + i + 1, 1)
+                i -= 1
+            elif isinstance(part, int):
+                del shape[i]
+            elif isinstance(part, slice):
+                start, stop, step = parse_slice(part, shape[i])
+                shape[i] = max(0, (stop - start + step - 1) // step)
+                i -= 1
+            else:
+                raise ValueError(part)
+
+        return Array[dtype, tuple(shape)]
+
+    if isinstance(domain, ProductDomain):
+        if isinstance(index, tuple):
+            assert len(index) == 1
+            index = index[0]
+        if isinstance(index, int):
+            return domain.__args__[index]
+        elif isinstance(index, slice):
+            return Product[domain.__args__[index]]
+        else:
+            raise ValueError(index)
+
+    raise NotImplementedError("TODO")
+
+
 @find_domain.register(ops.BinaryOp)
 def _find_domain_pointwise_binary_generic(op, lhs, rhs):
     if (

funsor/gaussian.py

@@ -445,7 +445,8 @@ def _eager_subs_real(self, subs, remaining_subs):
                 assert value.shape[-1] == self.inputs[k].num_elements
             values[k] = ops.expand(value, batch_shape + value.shape[-1:])
 
-        # Try to perform a complete substitution of all real variables, resulting in a Tensor.
+        # Try to perform a complete substitution of all real variables,
+        # resulting in a Tensor.
         if all(k in subs for k, d in self.inputs.items() if d.dtype == "real"):
             # Form the concatenated value.
             value = BlockVector(batch_shape + (event_size,))
@@ -461,8 +462,9 @@ def _eager_subs_real(self, subs, remaining_subs):
             assert result.output == Real
             return Subs(result, remaining_subs) if remaining_subs else result
 
-        # Perform a partial substution of a subset of real variables, resulting in a Joint.
-        # We split real inputs into two sets: a for the preserved and b for the substituted.
+        # Perform a partial substution of a subset of real variables, resulting
+        # in a Joint. We split real inputs into two sets: a for the preserved
+        # and b for the substituted.
         b = frozenset(k for k, v in subs.items())
         a = frozenset(
             k for k, d in self.inputs.items() if d.dtype == "real" and k not in b
@@ -703,7 +705,9 @@ def _sample(self, sampled_vars, sample_inputs, rng_key):
         sample_inputs = OrderedDict(
             (k, d) for k, d in sample_inputs.items() if k not in self.inputs
         )
-        sample_shape = tuple(int(d.dtype) for d in sample_inputs.values())
+        sample_shape = tuple(
+            d.size for d in sample_inputs.values() if d.dtype != "real"
+        )
         int_inputs = OrderedDict(
             (k, d) for k, d in self.inputs.items() if d.dtype != "real"
         )
@@ -716,32 +720,49 @@ def _sample(self, sampled_vars, sample_inputs, rng_key):
         if sampled_vars == frozenset(real_inputs):
             shape = sample_shape + self.info_vec.shape
             backend = get_backend()
-            if backend != "numpy":
+            info_vec = self.info_vec
+            precision_chol = self._precision_chol
+            if (
+                len(sample_inputs) == 1
+                and next(iter(sample_inputs.values())).dtype == "real"
+            ):
+                # Lazily compute a sample as a function of white noise.
+                for k, d in sample_inputs.items():
+                    white_noise = Variable(k, d)[tuple(int_inputs)]
+                info_vec = Tensor(info_vec, int_inputs)
+                precision_chol = Tensor(precision_chol, int_inputs)
+            elif backend == "numpy":
+                # Eagerly draw noise.
+                white_noise = np.random.randn(*shape)
+            else:
+                # Eagerly draw noise.
                 from importlib import import_module
 
                 dist = import_module(
                     funsor.distribution.BACKEND_TO_DISTRIBUTIONS_BACKEND[backend]
                 )
                 sample_args = (shape,) if rng_key is None else (rng_key, shape)
                 white_noise = dist.Normal.dist_class(0, 1).sample(*sample_args)
-            else:
-                white_noise = np.random.randn(*shape)
-            white_noise = ops.unsqueeze(white_noise, -1)
 
-            white_vec = ops.triangular_solve(
-                self.info_vec[..., None], self._precision_chol
-            )
+            # Jointly sample.
+            # This section may involve either Funsors or backend arrays.
+            white_vec = ops.triangular_solve(info_vec[..., None], precision_chol)
             sample = ops.triangular_solve(
-                white_noise + white_vec, self._precision_chol, transpose=True
+                white_noise[..., None] + white_vec, precision_chol, transpose=True
             )[..., 0]
+
+            # Extract shaped components.
             offsets, _ = _compute_offsets(real_inputs)
             results = []
             for key, domain in real_inputs.items():
-                data = sample[..., offsets[key] : offsets[key] + domain.num_elements]
-                data = data.reshape(shape[:-1] + domain.shape)
-                point = Tensor(data, inputs)
+                # TODO Support nontrivial slices in Funsor.__getitem__().
+                point = sample[..., offsets[key] : offsets[key] + domain.num_elements]
+                point = point.reshape(point.shape[:-1] + domain.shape)
+                if not isinstance(point, Funsor):  # I.e. when eagerly sampling.
+                    point = Tensor(point, inputs)
                 assert point.output == domain
                 results.append(Delta(key, point))
+
             results.append(self.log_normalizer)
             return reduce(ops.add, results)
 

funsor/montecarlo.py

@@ -6,10 +6,11 @@
 
 from funsor.cnf import Contraction
 from funsor.delta import Delta
+from funsor.gaussian import Gaussian
 from funsor.integrate import Integrate
 from funsor.interpretations import StatefulInterpretation
 from funsor.tensor import Tensor
-from funsor.terms import Approximate, Funsor, Number
+from funsor.terms import Approximate, Funsor, Number, Subs
 from funsor.util import get_backend
 
 from . import ops
@@ -86,8 +87,10 @@ def _extract_samples_contraction(discrete_density):
     return result
 
 
+@extract_samples.register(Subs)
 @extract_samples.register(Number)
 @extract_samples.register(Tensor)
+@extract_samples.register(Gaussian)
 def _extract_samples_scale(discrete_density):
     return {}
 

funsor/ops/builtin.py

@@ -13,6 +13,7 @@
     UNITS,
     BinaryOp,
     Op,
+    OpMeta,
     TransformOp,
     UnaryOp,
     declare_op_types,
@@ -43,6 +44,105 @@ def getitem(lhs, rhs, offset=0):
     return lhs[(slice(None),) * offset + (rhs,)]
 
 
+class GetsliceMeta(OpMeta):
+    """
+    Works around slice objects not being hashable.
+    """
+
+    def hash_args_kwargs(cls, args, kwargs):
+        index = args[0] if args else kwargs["index"]
+        if not isinstance(index, tuple):
+            index = (index,)
+        key = tuple(
+            (x.start, x.stop, x.step) if isinstance(x, slice) else x for x in index
+        )
+        return key
+
+
+@UnaryOp.make(metaclass=GetsliceMeta)
+def getslice(x, index=Ellipsis):
+    return x[index]
+
+
+getslice.supported_types = (type(None), type(Ellipsis), int, slice)
+
+
+def parse_ellipsis(index):
+    """
+    Helper to split a slice into parts left and right of Ellipses.
+
+    :param index: A tuple, or other object (None, int, slice, Funsor).
+    :returns: a pair of tuples ``left, right``.
+    :rtype: tuple
+    """
+    if not isinstance(index, tuple):
+        index = (index,)
+    left = []
+    i = 0
+    for part in index:
+        i += 1
+        if part is Ellipsis:
+            break
+        left.append(part)
+    right = []
+    for part in reversed(index[i:]):
+        if part is Ellipsis:
+            break
+        right.append(part)
+    right.reverse()
+    return tuple(left), tuple(right)
+
+
+def normalize_ellipsis(index, size):
+    """
+    Expand Ellipses in an index to fill the given number of dimensions.
+
+    This should satisfy the equation::
+
+        x[i] == x[normalize_ellipsis(i, len(x.shape))]
+    """
+    left, right = parse_ellipsis(index)
+    if len(left) + len(right) > size:
+        raise ValueError(f"Index is too wide: {index}")
+    middle = (slice(None),) * (size - len(left) - len(right))
+    return left + middle + right
+
+
+def parse_slice(s, size):
+    """
+    Helper to determine nonnegative integers (start, stop, step) of a slice.
+
+    :param slice s: A slice.
+    :param int size: The size of the array being indexed into.
+    :returns: A tuple of nonnegative integers ``start, stop, step``.
+    :rtype: tuple
+    """
+    start = s.start
+    if start is None:
+        start = 0
+    assert isinstance(start, int)
+    if start >= 0:
+        start = min(size, start)
+    else:
+        start = max(0, size + start)
+
+    stop = s.stop
+    if stop is None:
+        stop = size
+    assert isinstance(stop, int)
+    if stop >= 0:
+        stop = min(size, stop)
+    else:
+        stop = max(0, size + stop)
+
+    step = s.step
+    if step is None:
+        step = 1
+    assert isinstance(step, int)
+
+    return start, stop, step
+
+
 abs = UnaryOp.make(_builtin_abs)
 eq = BinaryOp.make(operator.eq)
 ge = BinaryOp.make(operator.ge)
@@ -194,6 +294,7 @@ def sigmoid_log_abs_det_jacobian(x, y):
     "floordiv",
     "ge",
     "getitem",
+    "getslice",
     "gt",
     "invert",
     "le",