Repository: pyro-ppl/funsor

Issue:
Implement Monte Carlo interpretation of .logaddexp()
It appears that the natural way to express Monte Carlo sampling in funsor is with a nonstandard interpretation of `.logaddexp()` reductions. (Note that for numerical stability, we represent densities in log-space; hence expectation is represented by `.logaddexp()` reductions rather than `.sum()` reductions).

To perform sampling on factor graphs, we will orthogonally need #50.

## Tasks

- [ ] #75 Implement monte carlo `.sample()` methods
- [ ] Implement a funsor `Integrate(log_measure, integrand, reduced_vars)`
- [ ] Use `Integrate()` in test_samplers.py to test `Gaussian` moments
- [ ] Implement a `monte_carlo` interpretation that uses samplers in `Integrate()`

Discussion:
fritzo: @eb8680 This is now ready for review. It's probably easier for me to walk you through the changes, rather than split this into smaller PRs. I've split examples/vae.py out as #95 since that involves shape issues.

Patch:
diff --git a/funsor/__init__.py b/funsor/__init__.py
index 7b9b9fe..2e70fa7 100644
--- a/funsor/__init__.py
+++ b/funsor/__init__.py
@@ -1,6 +1,7 @@
 from __future__ import absolute_import, division, print_function
 
 from funsor.domains import Domain, bint, find_domain, reals
+from funsor.integrate import Integrate
 from funsor.interpreter import reinterpret
 from funsor.terms import Funsor, Number, Variable, of_shape, to_data, to_funsor
 from funsor.torch import Tensor, arange, torch_einsum
@@ -14,9 +15,11 @@ from . import (
     einsum,
     gaussian,
     handlers,
+    integrate,
     interpreter,
     joint,
     minipyro,
+    montecarlo,
     ops,
     sum_product,
     terms,
@@ -26,14 +29,15 @@ from . import (
 __all__ = [
     'Domain',
     'Funsor',
+    'Integrate',
     'Number',
     'Tensor',
     'Variable',
     'adjoint',
     'arange',
     'backward',
-    'contract',
     'bint',
+    'contract',
     'delta',
     'distributions',
     'domains',
@@ -41,9 +45,11 @@ __all__ = [
     'find_domain',
     'gaussian',
     'handlers',
+    'integrate',
     'interpreter',
     'joint',
     'minipyro',
+    'montecarlo',
     'of_shape',
     'ops',
     'reals',
diff --git a/funsor/delta.py b/funsor/delta.py
index e04527e..9f9880d 100644
--- a/funsor/delta.py
+++ b/funsor/delta.py
@@ -6,6 +6,7 @@ from six import add_metaclass
 
 import funsor.ops as ops
 from funsor.domains import reals
+from funsor.integrate import Integrate, integrator
 from funsor.ops import Op
 from funsor.terms import Align, Binary, Funsor, FunsorMeta, Number, Variable, eager, to_funsor
 
@@ -107,6 +108,16 @@ def eager_binary(op, lhs, rhs):
     return None  # defer to default implementation
 
 
+@eager.register(Integrate, Delta, Funsor, frozenset)
+@integrator
+def eager_integrate(delta, integrand, reduced_vars):
+    assert delta.name in reduced_vars
+    integrand = integrand.eager_subs(((delta.name, delta.point),))
+    log_measure = delta.log_density
+    reduced_vars -= frozenset([delta.name])
+    return Integrate(log_measure, integrand, reduced_vars)
+
+
 __all__ = [
     'Delta',
 ]
diff --git a/funsor/gaussian.py b/funsor/gaussian.py
index b6add58..82b5d18 100644
--- a/funsor/gaussian.py
+++ b/funsor/gaussian.py
@@ -11,8 +11,10 @@ from six.moves import reduce
 import funsor.ops as ops
 from funsor.delta import Delta
 from funsor.domains import reals
+from funsor.integrate import Integrate, integrator
+from funsor.montecarlo import monte_carlo
 from funsor.ops import AddOp
-from funsor.terms import Binary, Funsor, FunsorMeta, Number, eager
+from funsor.terms import Binary, Funsor, FunsorMeta, Number, Variable, eager
 from funsor.torch import Tensor, align_tensor, align_tensors, materialize
 from funsor.util import lazy_property
 
@@ -26,10 +28,14 @@ def _issubshape(subshape, supershape):
     return True
 
 
-def _log_det_tril(x):
+def _log_det_tri(x):
     return x.diagonal(dim1=-1, dim2=-2).log().sum(-1)
 
 
+def _det_tri(x):
+    return x.diagonal(dim1=-1, dim2=-2).prod(-1)
+
+
 def _mv(mat, vec):
     return torch.matmul(mat, vec.unsqueeze(-1)).squeeze(-1)
 
@@ -44,6 +50,16 @@ def _vmv(mat, vec):
     return result.squeeze(-1).squeeze(-1)
 
 
+def _trace_mm(x, y):
+    """
+    Computes ``trace(x @ y)``.
+    """
+    assert x.dim() >= 2
+    assert y.dim() >= 2
+    xy = x * y
+    return xy.reshape(xy.shape[:-2] + (-1,)).sum(-1)
+
+
 def _compute_offsets(inputs):
     """
     Compute offsets of real inputs into the concatenated Gaussian dims.
@@ -195,7 +211,7 @@ class Gaussian(Funsor):
                        Tensor(self.precision, int_inputs)]
             tensors.extend(subs.values())
             inputs, tensors = align_tensors(*tensors)
-            batch_dim = self.loc.dim() - 1
+            batch_dim = tensors[0].dim() - 1
             batch_shape = broadcast_shape(*(x.shape[:batch_dim] for x in tensors))
             (loc, precision), values = tensors[:2], tensors[2:]
 
@@ -219,8 +235,8 @@ class Gaussian(Funsor):
     @lazy_property
     def _log_normalizer(self):
         dim = self.loc.size(-1)
-        log_det_term = _log_det_tril(torch.cholesky(self.precision))
-        data = log_det_term - 0.5 * math.log(2 * math.pi) * dim
+        log_det_term = _log_det_tri(torch.cholesky(self.precision))
+        data = -log_det_term + 0.5 * math.log(2 * math.pi) * dim
         inputs = OrderedDict((k, v) for k, v in self.inputs.items() if v.dtype != 'real')
         return Tensor(data, inputs)
 
@@ -247,15 +263,15 @@ class Gaussian(Funsor):
                 index = torch.tensor(index)
 
                 loc = self.loc[..., index]
-                self_scale_tril = torch.inverse(torch.cholesky(self.precision))
-                self_covariance = torch.matmul(self_scale_tril, self_scale_tril.transpose(-1, -2))
+                self_scale_tri = torch.inverse(torch.cholesky(self.precision)).transpose(-1, -2)
+                self_covariance = torch.matmul(self_scale_tri, self_scale_tri.transpose(-1, -2))
                 covariance = self_covariance[..., index.unsqueeze(-1), index]
-                scale_tril = torch.cholesky(covariance)
-                inv_scale_tril = torch.inverse(scale_tril)
-                precision = torch.matmul(inv_scale_tril, inv_scale_tril.transpose(-1, -2))
+                scale_tri = torch.cholesky(covariance)
+                inv_scale_tri = torch.inverse(scale_tri)
+                precision = torch.matmul(inv_scale_tri.transpose(-1, -2), inv_scale_tri)
                 reduced_dim = sum(self.inputs[k].num_elements for k in reduced_reals)
-                log_det_term = _log_det_tril(scale_tril) - _log_det_tril(self_scale_tril)
-                log_prob = Tensor(log_det_term - 0.5 * math.log(2 * math.pi) * reduced_dim, int_inputs)
+                log_det_term = _log_det_tri(self_scale_tri) - _log_det_tri(scale_tri)
+                log_prob = Tensor(log_det_term + 0.5 * math.log(2 * math.pi) * reduced_dim, int_inputs)
                 result = log_prob + Gaussian(loc, precision, inputs)
 
             return result.reduce(ops.logaddexp, reduced_ints)
@@ -265,7 +281,7 @@ class Gaussian(Funsor):
 
         return None  # defer to default implementation
 
-    def sample(self, sampled_vars, sample_inputs=None):
+    def unscaled_sample(self, sampled_vars, sample_inputs=None):
         sampled_vars = sampled_vars.intersection(self.inputs)
         if not sampled_vars:
             return self
@@ -274,7 +290,9 @@ class Gaussian(Funsor):
         # Partition inputs into sample_inputs + int_inputs + real_inputs.
         if sample_inputs is None:
             sample_inputs = OrderedDict()
-        assert frozenset(sample_inputs).isdisjoint(self.inputs)
+        else:
+            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())
         int_inputs = OrderedDict((k, d) for k, d in self.inputs.items() if d.dtype != 'real')
         real_inputs = OrderedDict((k, d) for k, d in self.inputs.items() if d.dtype == 'real')
@@ -282,11 +300,11 @@ class Gaussian(Funsor):
         inputs.update(int_inputs)
 
         if sampled_vars == frozenset(real_inputs):
-            scale_tril = torch.inverse(torch.cholesky(self.precision))
-            assert self.loc.shape == scale_tril.shape[:-1]
+            scale_tri = torch.inverse(torch.cholesky(self.precision)).transpose(-1, -2)
+            assert self.loc.shape == scale_tri.shape[:-1]
             shape = sample_shape + self.loc.shape
             white_noise = torch.randn(shape)
-            sample = self.loc + _mv(scale_tril, white_noise)
+            sample = self.loc + _mv(scale_tri, white_noise)
             offsets, _ = _compute_offsets(real_inputs)
             results = []
             for key, domain in real_inputs.items():
@@ -317,13 +335,71 @@ def eager_add_gaussian_gaussian(op, lhs, rhs):
     # Fuse aligned Gaussians.
     precision_loc = _mv(lhs_precision, lhs_loc) + _mv(rhs_precision, rhs_loc)
     precision = lhs_precision + rhs_precision
-    scale_tril = torch.inverse(torch.cholesky(precision))
-    loc = _mv(scale_tril.transpose(-1, -2), _mv(scale_tril, precision_loc))
+    scale_tri = torch.inverse(torch.cholesky(precision)).transpose(-1, -2)
+    loc = _mv(scale_tri, _mv(scale_tri.transpose(-1, -2), precision_loc))
     quadratic_term = _vmv(lhs_precision, loc - lhs_loc) + _vmv(rhs_precision, loc - rhs_loc)
     likelihood = Tensor(-0.5 * quadratic_term, int_inputs)
     return likelihood + Gaussian(loc, precision, inputs)
 
 
+@eager.register(Integrate, Gaussian, Variable, frozenset)
+@integrator
+def eager_integrate(log_measure, integrand, reduced_vars):
+    real_vars = frozenset(k for k in reduced_vars if log_measure.inputs[k].dtype == 'real')
+    if real_vars:
+        assert real_vars == frozenset([integrand.name])
+        data = log_measure.loc * log_measure._log_normalizer.data.exp().unsqueeze(-1)
+        data = data.reshape(log_measure.loc.shape[:-1] + integrand.output.shape)
+        inputs = OrderedDict((k, d) for k, d in log_measure.inputs.items() if d.dtype != 'real')
+        return Tensor(data, inputs)
+
+    return None  # defer to default implementation
+
+
+@eager.register(Integrate, Gaussian, Gaussian, frozenset)
+@integrator
+def eager_integrate(log_measure, integrand, reduced_vars):
+    real_vars = frozenset(k for k in reduced_vars if log_measure.inputs[k].dtype == 'real')
+    if real_vars:
+
+        lhs_reals = frozenset(k for k, d in log_measure.inputs.items() if d.dtype == 'real')
+        rhs_reals = frozenset(k for k, d in integrand.inputs.items() if d.dtype == 'real')
+        if lhs_reals == real_vars and rhs_reals <= real_vars:
+            inputs = OrderedDict((k, d) for t in (log_measure, integrand)
+                                 for k, d in t.inputs.items())
+            lhs_loc, lhs_precision = align_gaussian(inputs, log_measure)
+            rhs_loc, rhs_precision = align_gaussian(inputs, integrand)
+
+            # Compute the expectation of a non-normalized quadratic form.
+            # See "The Matrix Cookbook" (November 15, 2012) ss. 8.2.2 eq. 380.
+            # http://www.math.uwaterloo.ca/~hwolkowi/matrixcookbook.pdf
+            lhs_scale_tri = torch.inverse(torch.cholesky(lhs_precision)).transpose(-1, -2)
+            lhs_covariance = torch.matmul(lhs_scale_tri, lhs_scale_tri.transpose(-1, -2))
+            dim = lhs_loc.size(-1)
+            norm = _det_tri(lhs_scale_tri) * (2 * math.pi) ** (0.5 * dim)
+            data = -0.5 * norm * (_vmv(rhs_precision, lhs_loc - rhs_loc) +
+                                  _trace_mm(rhs_precision, lhs_covariance))
+            inputs = OrderedDict((k, d) for k, d in inputs.items() if k not in reduced_vars)
+            result = Tensor(data, inputs)
+            return result.reduce(ops.add, reduced_vars - real_vars)
+
+        raise NotImplementedError('TODO implement partial integration')
+
+    return None  # defer to default implementation
+
+
+@monte_carlo.register(Integrate, Gaussian, Funsor, frozenset)
+@integrator
+def monte_carlo_integrate(log_measure, integrand, reduced_vars):
+    real_vars = frozenset(k for k in reduced_vars if log_measure.inputs[k].dtype == 'real')
+    if real_vars:
+        log_measure = log_measure.sample(real_vars, monte_carlo.sample_inputs)
+        reduced_vars = reduced_vars | frozenset(monte_carlo.sample_inputs)
+        return Integrate(log_measure, integrand, reduced_vars)
+
+    return None  # defer to default implementation
+
+
 __all__ = [
     'Gaussian',
     'align_gaussian',
diff --git a/funsor/integrate.py b/funsor/integrate.py
new file mode 100644
index 0000000..7799b1a
--- /dev/null
+++ b/funsor/integrate.py
@@ -0,0 +1,85 @@
+from __future__ import absolute_import, division, print_function
+
+import functools
+from collections import OrderedDict
+
+import funsor.ops as ops
+from funsor.contract import Contract
+from funsor.terms import Funsor, Reduce, eager
+
+
+class Integrate(Funsor):
+    """
+    Funsor representing an integral wrt a log density funsor.
+    """
+    def __init__(self, log_measure, integrand, reduced_vars):
+        assert isinstance(log_measure, Funsor)
+        assert isinstance(integrand, Funsor)
+        assert isinstance(reduced_vars, frozenset)
+        inputs = OrderedDict((k, d) for term in (log_measure, integrand)
+                             for (k, d) in term.inputs.items()
+                             if k not in reduced_vars)
+        output = integrand.output
+        super(Integrate, self).__init__(inputs, output)
+        self.log_measure = log_measure
+        self.integrand = integrand
+        self.reduced_vars = reduced_vars
+
+    def eager_subs(self, subs):
+        raise NotImplementedError('TODO')
+
+
+def _simplify_integrate(fn, log_measure, integrand, reduced_vars):
+    """
+    Reduce free variables that do not appear in both inputs.
+    """
+    if not reduced_vars:
+        return log_measure.exp() * integrand
+
+    log_measure_vars = frozenset(log_measure.inputs)
+    integrand_vars = frozenset(integrand.inputs)
+    assert reduced_vars <= log_measure_vars | integrand_vars
+    progress = False
+    if not reduced_vars <= log_measure_vars:
+        integrand = integrand.reduce(ops.add, reduced_vars - log_measure_vars)
+        reduced_vars = reduced_vars & log_measure_vars
+        progress = True
+    if not reduced_vars <= integrand_vars:
+        log_measure = log_measure.reduce(ops.logaddexp, reduced_vars - integrand_vars)
+        reduced_vars = reduced_vars & integrand_vars
+        progress = True
+    if progress:
+        return Integrate(log_measure, integrand, reduced_vars)
+
+    return fn(log_measure, integrand, reduced_vars)
+
+
+def integrator(fn):
+    """
+    Decorator for integration implementations.
+    """
+    return functools.partial(_simplify_integrate, fn)
+
+
+@eager.register(Integrate, Funsor, Funsor, frozenset)
+@integrator
+def eager_integrate(log_measure, integrand, reduced_vars):
+    return Contract(log_measure.exp(), integrand, reduced_vars)
+
+
+@eager.register(Integrate, Reduce, Funsor, frozenset)
+@integrator
+def eager_integrate(log_measure, integrand, reduced_vars):
+    if log_measure.op is ops.logaddexp:
+        if not log_measure.reduced_vars.isdisjoint(reduced_vars):
+            raise NotImplementedError('TODO alpha convert')
+        arg = Integrate(log_measure.arg, integrand, reduced_vars)
+        return arg.reduce(ops.add, log_measure.reduced_vars)
+
+    return Contract(log_measure.exp(), integrand, reduced_vars)
+
+
+__all__ = [
+    'Integrate',
+    'integrator',
+]
diff --git a/funsor/interpreter.py b/funsor/interpreter.py
index 9904fa2..bab0205 100644
--- a/funsor/interpreter.py
+++ b/funsor/interpreter.py
@@ -1,6 +1,7 @@
 from __future__ import absolute_import, division, print_function
 
 import os
+import re
 import types
 from collections import OrderedDict
 
@@ -24,12 +25,18 @@ if _DEBUG:
         indent = '  ' * _STACK_SIZE
         typenames = [cls.__name__] + [type(arg).__name__ for arg in args]
         print(indent + ' '.join(typenames))
+
         _STACK_SIZE += 1
         try:
             result = _INTERPRETATION(cls, *args)
         finally:
             _STACK_SIZE -= 1
-        print(indent + '-> ' + type(result).__name__)
+
+        if _DEBUG > 1:
+            result_str = re.sub('\n', '\n          ' + indent, str(result))
+        else:
+            result_str = type(result).__name__
+        print(indent + '-> ' + result_str)
         return result
 else:
     def interpret(cls, *args):
diff --git a/funsor/joint.py b/funsor/joint.py
index 335b870..4108e83 100644
--- a/funsor/joint.py
+++ b/funsor/joint.py
@@ -1,5 +1,6 @@
 from __future__ import absolute_import, division, print_function
 
+import functools
 from collections import OrderedDict
 
 from six import add_metaclass
@@ -9,9 +10,11 @@ import funsor.ops as ops
 from funsor.delta import Delta
 from funsor.domains import reals
 from funsor.gaussian import Gaussian
+from funsor.integrate import Integrate, integrator
+from funsor.montecarlo import monte_carlo
 from funsor.ops import AddOp, Op
-from funsor.terms import Align, Binary, Funsor, FunsorMeta, Number, eager, to_funsor
-from funsor.torch import Tensor
+from funsor.terms import Align, Binary, Funsor, FunsorMeta, Number, Variable, eager, to_funsor
+from funsor.torch import Tensor, arange
 
 
 class JointMeta(FunsorMeta):
@@ -123,22 +126,17 @@ class Joint(Funsor):
 
         return None  # defer to default implementation
 
-    def sample(self, sampled_vars, sample_inputs=None):
-        if sample_inputs is None:
-            sample_inputs = OrderedDict()
-        assert frozenset(sample_inputs).isdisjoint(self.inputs)
+    def unscaled_sample(self, sampled_vars, sample_inputs=None):
         discrete_vars = sampled_vars.intersection(self.discrete.inputs)
         gaussian_vars = frozenset(k for k in sampled_vars
                                   if k in self.gaussian.inputs
                                   if self.gaussian.inputs[k].dtype == 'real')
         result = self
         if discrete_vars:
-            discrete = result.discrete.sample(discrete_vars, sample_inputs)
+            discrete = result.discrete.unscaled_sample(discrete_vars, sample_inputs)
             result = Joint(result.deltas, gaussian=result.gaussian) + discrete
         if gaussian_vars:
-            sample_inputs = OrderedDict((k, v) for k, v in sample_inputs.items()
-                                        if k not in result.gaussian.inputs)
-            gaussian = result.gaussian.sample(gaussian_vars, sample_inputs)
+            gaussian = result.gaussian.unscaled_sample(gaussian_vars, sample_inputs)
             result = Joint(result.deltas, result.discrete) + gaussian
         return result
 
@@ -275,6 +273,86 @@ def eager_binary(op, gaussian, discrete):
         return Joint(discrete=-discrete, gaussian=gaussian)
 
 
+################################################################################
+# Patterns for integration
+################################################################################
+
+def _simplify_integrate(fn, joint, integrand, reduced_vars):
+    if any(d.name in reduced_vars for d in joint.deltas):
+        subs = tuple((d.name, d.point) for d in joint.deltas if d.name in reduced_vars)
+        deltas = tuple(d for d in joint.deltas if d.name not in reduced_vars)
+        log_measure = Joint(deltas, joint.discrete, joint.gaussian)
+        integrand = integrand.eager_subs(subs)
+        reduced_vars = reduced_vars - frozenset(name for name, point in subs)
+        return Integrate(log_measure, integrand, reduced_vars)
+
+    return fn(log_measure, integrand, reduced_vars)
+
+
+def _joint_integrator(fn):
+    """
+    Decorator for Integrate(Joint(...), ...) patterns.
+    """
+    return integrator(functools.partial(_simplify_integrate, fn))
+
+
+@eager.register(Integrate, Joint, Delta, frozenset)
+@_joint_integrator
+def eager_integrate(log_measure, integrand, reduced_vars):
+    raise NotImplementedError('TODO')
+
+
+@eager.register(Integrate, Joint, Tensor, frozenset)
+@_joint_integrator
+def eager_integrate(log_measure, integrand, reduced_vars):
+    raise NotImplementedError('TODO')
+
+
+@eager.register(Integrate, Joint, Gaussian, frozenset)
+@_joint_integrator
+def eager_integrate(log_measure, integrand, reduced_vars):
+    raise NotImplementedError('TODO')
+
+
+@eager.register(Integrate, Joint, Joint, frozenset)
+@_joint_integrator
+def eager_integrate(log_measure, integrand, reduced_vars):
+    raise NotImplementedError('TODO')
+
+
+@eager.register(Integrate, Joint, Variable, frozenset)
+@integrator
+def eager_integrate(log_measure, integrand, reduced_vars):
+    name = integrand.name
+    assert reduced_vars == frozenset([name])
+    if any(d.name == name for d in log_measure.deltas):
+        deltas = tuple(d for d in log_measure.deltas if d.name != name)
+        log_norm = Joint(deltas, log_measure.discrete, log_measure.gaussian)
+        for d in log_measure.deltas:
+            if d.name == name:
+                mean = d.point
+                break
+        return mean * log_norm.exp()
+    elif name in log_measure.discrete.inputs:
+        integrand = arange(name, integrand.inputs[name].dtype)
+        return Integrate(log_measure, integrand, reduced_vars)
+    else:
+        assert name in log_measure.gaussian.inputs
+        gaussian = Integrate(log_measure.gaussian, integrand, reduced_vars)
+        return Joint(log_measure.deltas, log_measure.discrete).exp() * gaussian
+
+
+@monte_carlo.register(Integrate, Joint, Funsor, frozenset)
+@integrator
+def monte_carlo_integrate(log_measure, integrand, reduced_vars):
+    sampled_log_measure = log_measure.sample(reduced_vars, monte_carlo.sample_inputs)
+    if sampled_log_measure is not log_measure:
+        reduced_vars = reduced_vars | frozenset(monte_carlo.sample_inputs)
+        return Integrate(sampled_log_measure, integrand, reduced_vars)
+
+    return None  # defer to default implementation
+
+
 __all__ = [
     'Joint',
 ]
diff --git a/funsor/montecarlo.py b/funsor/montecarlo.py
new file mode 100644
index 0000000..8e7c69f
--- /dev/null
+++ b/funsor/montecarlo.py
@@ -0,0 +1,47 @@
+from __future__ import absolute_import, division, print_function
+
+from collections import OrderedDict
+
+from contextlib2 import contextmanager
+
+from funsor.interpreter import dispatched_interpretation, interpretation
+from funsor.terms import eager
+
+
+@dispatched_interpretation
+def monte_carlo(cls, *args):
+    """
+    A Monte Carlo interpretation of :class:`~funsor.integrate.Integrate`
+    expressions. This falls back to :class:`~funsor.terms.eager` in other
+    cases.
+    """
+    # TODO Memoize sample statements in a context manager.
+    result = monte_carlo.dispatch(cls, *args)
+    if result is None:
+        result = eager(cls, *args)
+    return result
+
+
+# This is a globally configurable parameter to draw multiple samples.
+monte_carlo.sample_inputs = OrderedDict()
+
+
+@contextmanager
+def monte_carlo_interpretation(**sample_inputs):
+    """
+    Context manager to set ``monte_carlo.sample_inputs`` and
+    install the :func:`monte_carlo` interpretation.
+    """
+    old = monte_carlo.sample_inputs
+    monte_carlo.sample_inputs = OrderedDict(sample_inputs)
+    try:
+        with interpretation(monte_carlo):
+            yield
+    finally:
+        monte_carlo.sample_inputs = old
+
+
+__all__ = [
+    'monte_carlo',
+    'monte_carlo_interpretation'
+]
diff --git a/funsor/ops.py b/funsor/ops.py
index dae6b9e..61c7feb 100644
--- a/funsor/ops.py
+++ b/funsor/ops.py
@@ -58,8 +58,8 @@ class GetitemOp(Op):
         assert offset >= 0
         self.offset = offset
         self._prefix = (slice(None),) * offset
-        self.__name__ = 'GetitemOp({})'.format(offset)
         super(GetitemOp, self).__init__(self._default)
+        self.__name__ = 'GetitemOp({})'.format(offset)
 
     def _default(self, x, y):
         return x[self._prefix + (y,)] if self.offset else x[y]
diff --git a/funsor/terms.py b/funsor/terms.py
index aada944..0ad94f2 100644
--- a/funsor/terms.py
+++ b/funsor/terms.py
@@ -2,6 +2,7 @@ from __future__ import absolute_import, division, print_function
 
 import functools
 import itertools
+import math
 import numbers
 from abc import ABCMeta, abstractmethod
 from collections import Hashable, OrderedDict
@@ -192,6 +193,15 @@ class Funsor(object):
     def __nonzero__(self):
         return self.__bool__()
 
+    def __len__(self):
+        if not self.output.shape:
+            raise ValueError('Funsor with empty shape has no len()')
+        return self.output.shape[0]
+
+    def __iter__(self):
+        for i in range(len(self)):
+            yield self[i]
+
     def item(self):
         if self.inputs or self.output.shape:
             raise ValueError(
@@ -225,9 +235,10 @@ class Funsor(object):
         Create a Monte Carlo approximation to this funsor by replacing
         functions of ``sampled_vars`` with :class:`~funsor.delta.Delta`s.
 
-        If ``sample_inputs`` is not provided, the result is a :class:`Funsor`
-        with the same ``.inputs`` and ``.output`` as the original funsor, so
-        that self can be replaced by the sample in expectation computations::
+        The result is a :class:`Funsor` with the same ``.inputs`` and
+        ``.output`` as the original funsor (plus ``sample_inputs`` if
+        provided), so that self can be replaced by the sample in expectation
+        computations::
 
             y = x.sample(sampled_vars)
             assert y.inputs == x.inputs
@@ -236,13 +247,7 @@ class Funsor(object):
             approx = (y.exp() * integrand).reduce(ops.add)
 
         If ``sample_inputs`` is provided, this creates a batch of samples
-        that are intended to be averaged, however this reduction is not
-        performed by the :meth:`sample` method::
-
-            y = x.sample(sampled_vars, sample_inputs)
-            total = reduce(ops.mul, d.num_elements) for d in sample_inputs.values())
-            exact = (x.exp() * integrand).reduce(ops.add)
-            approx = (y.exp() * integrand).reduce(ops.add) / total
+        scaled samples.
 
         :param frozenset sampled_vars: A set of input variables to sample.
         :param OrderedDict sample_inputs: An optional mapping from variable
@@ -251,6 +256,25 @@ class Funsor(object):
         """
         assert self.output == reals()
         assert isinstance(sampled_vars, frozenset)
+        if sampled_vars.isdisjoint(self.inputs):
+            return self
+
+        result = self.unscaled_sample(sampled_vars, sample_inputs)
+        if sample_inputs is not None:
+            log_scale = 0
+            for var, domain in sample_inputs.items():
+                if var in result.inputs and var not in self.inputs:
+                    log_scale -= math.log(domain.dtype)
+            if log_scale != 0:
+                result += log_scale
+        return result
+
+    def unscaled_sample(self, sampled_vars, sample_inputs=None):
+        """
+        Internal method to draw an unscaled sample.
+        """
+        assert self.output == reals()
+        assert isinstance(sampled_vars, frozenset)
         if sampled_vars.isdisjoint(self.inputs):
             return self
         raise NotImplementedError
diff --git a/funsor/torch.py b/funsor/torch.py
index afd4a27..897a94c 100644
--- a/funsor/torch.py
+++ b/funsor/torch.py
@@ -13,6 +13,8 @@ import funsor.ops as ops
 from funsor.contract import Contract, contractor
 from funsor.delta import Delta
 from funsor.domains import Domain, bint, find_domain, reals
+from funsor.integrate import Integrate, integrator
+from funsor.montecarlo import monte_carlo
 from funsor.ops import GetitemOp, Op
 from funsor.six import getargspec
 from funsor.terms import Binary, Funsor, FunsorMeta, Number, Variable, eager, to_data, to_funsor
@@ -216,7 +218,7 @@ class Tensor(Funsor):
             assert isinstance(reduced_vars, frozenset)
             self_vars = frozenset(self.inputs)
             reduced_vars = reduced_vars & self_vars
-            if reduced_vars == self_vars:
+            if reduced_vars == self_vars and not self.output.shape:
                 # Reduce all dims at once.
                 if op is ops.logaddexp:
                     # work around missing torch.Tensor.logsumexp()
@@ -240,7 +242,7 @@ class Tensor(Funsor):
             return Tensor(data, inputs, self.dtype)
         return super(Tensor, self).eager_reduce(op, reduced_vars)
 
-    def sample(self, sampled_vars, sample_inputs=None):
+    def unscaled_sample(self, sampled_vars, sample_inputs=None):
         assert self.output == reals()
         sampled_vars = sampled_vars.intersection(self.inputs)
         if not sampled_vars:
@@ -249,7 +251,9 @@ class Tensor(Funsor):
         # Partition inputs into sample_inputs + batch_inputs + event_inputs.
         if sample_inputs is None:
             sample_inputs = OrderedDict()
-        assert frozenset(sample_inputs).isdisjoint(self.inputs)
+        else:
+            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())
         batch_inputs = OrderedDict((k, d) for k, d in self.inputs.items() if k not in sampled_vars)
         event_inputs = OrderedDict((k, d) for k, d in self.inputs.items() if k in sampled_vars)
@@ -342,6 +346,21 @@ def eager_binary_tensor_tensor(op, lhs, rhs):
     else:
         inputs, (lhs_data, rhs_data) = align_tensors(lhs, rhs)
 
+    # Reshape to support broadcasting of output shape.
+    if inputs:
+        lhs_dim = len(lhs.output.shape)
+        rhs_dim = len(rhs.output.shape)
+        if lhs_dim < rhs_dim:
+            cut = lhs_data.dim() - lhs_dim
+            shape = lhs_data.shape
+            shape = shape[:cut] + (1,) * (rhs_dim - lhs_dim) + shape[cut:]
+            lhs_data = lhs_data.reshape(shape)
+        elif rhs_dim < lhs_dim:
+            cut = rhs_data.dim() - rhs_dim
+            shape = rhs_data.shape
+            shape = shape[:cut] + (1,) * (lhs_dim - rhs_dim) + shape[cut:]
+            rhs_data = rhs_data.reshape(shape)
+
     data = op(lhs_data, rhs_data)
     return Tensor(data, inputs, dtype)
 
@@ -412,6 +431,20 @@ def eager_contract(lhs, rhs, reduced_vars):
     return Tensor(data, inputs, dtype)
 
 
+@monte_carlo.register(Integrate, Tensor, Tensor, frozenset)
+@integrator
+def eager_integrate(log_measure, integrand, reduced_vars):
+    return Contract(log_measure.exp(), integrand, reduced_vars)
+
+
+@monte_carlo.register(Integrate, Tensor, Funsor, frozenset)
+@integrator
+def monte_carlo_integrate(log_measure, integrand, reduced_vars):
+    log_measure = log_measure.sample(reduced_vars, monte_carlo.sample_inputs)
+    reduced_vars = reduced_vars | frozenset(monte_carlo.sample_inputs)
+    return Integrate(log_measure, integrand, reduced_vars)
+
+
 def arange(name, size):
     """
     Helper to create a named :func:`torch.arange` funsor.
@@ -538,7 +571,10 @@ class _Memoized(object):
 
 
 def _function(inputs, output, fn):
-    names = getargspec(fn)[0]
+    if isinstance(fn, torch.nn.Module):
+        names = getargspec(fn.forward)[0][1:]
+    else:
+        names = getargspec(fn)[0]
     args = tuple(Variable(name, domain) for (name, domain) in zip(names, inputs))
     assert len(args) == len(inputs)
     if not isinstance(output, Domain):


Test patch:
diff --git a/funsor/testing.py b/funsor/testing.py
index 3c05d28..1e6e640 100644
--- a/funsor/testing.py
+++ b/funsor/testing.py
@@ -78,16 +78,16 @@ def assert_close(actual, expected, atol=1e-6, rtol=1e-6):
                 actual = actual[~eq]
                 expected = expected[~eq]
             diff = (actual.detach() - expected.detach()).abs()
-            if atol is not None:
-                assert diff.max() < atol, msg
             if rtol is not None:
                 assert (diff / (atol + expected.detach().abs())).max() < rtol, msg
+            elif atol is not None:
+                assert diff.max() < atol, msg
     elif isinstance(actual, numbers.Number):
         diff = abs(actual - expected)
-        if atol is not None:
-            assert diff < atol, msg
         if rtol is not None:
             assert diff < (atol + expected) * rtol, msg
+        elif atol is not None:
+            assert diff < atol, msg
     else:
         raise ValueError('cannot compare objects of type {}'.format(type(actual)))
 
@@ -186,6 +186,7 @@ def random_gaussian(inputs):
     loc = torch.randn(batch_shape + event_shape)
     prec_sqrt = torch.randn(batch_shape + event_shape + event_shape)
     precision = torch.matmul(prec_sqrt, prec_sqrt.transpose(-1, -2))
+    precision = precision + 0.05 * torch.eye(event_shape[0])
     return Gaussian(loc, precision, inputs)
 
 
diff --git a/test/test_gaussian.py b/test/test_gaussian.py
index 699709f..d0dcba0 100644
--- a/test/test_gaussian.py
+++ b/test/test_gaussian.py
@@ -5,12 +5,15 @@ from collections import OrderedDict
 
 import pytest
 import torch
+from six.moves import reduce
 
 import funsor.ops as ops
 from funsor.domains import bint, reals
 from funsor.gaussian import Gaussian
+from funsor.integrate import Integrate
 from funsor.joint import Joint
-from funsor.terms import Number
+from funsor.montecarlo import monte_carlo_interpretation
+from funsor.terms import Number, Variable
 from funsor.testing import assert_close, id_from_inputs, random_gaussian, random_tensor, xfail_if_not_implemented
 from funsor.torch import Tensor
 
@@ -225,3 +228,62 @@ def test_logsumexp(int_inputs, real_inputs):
     g_xy = g.reduce(ops.logaddexp, frozenset(['x', 'y']))
     assert_close(g_xy, g.reduce(ops.logaddexp, 'x').reduce(ops.logaddexp, 'y'), atol=1e-3, rtol=None)
     assert_close(g_xy, g.reduce(ops.logaddexp, 'y').reduce(ops.logaddexp, 'x'), atol=1e-3, rtol=None)
+
+
+@pytest.mark.parametrize('int_inputs', [
+    {},
+    {'i': bint(2)},
+], ids=id_from_inputs)
+@pytest.mark.parametrize('real_inputs', [
+    {'x': reals()},
+    {'x': reals(4)},
+    {'x': reals(2, 3)},
+], ids=id_from_inputs)
+def test_integrate_variable(int_inputs, real_inputs):
+    int_inputs = OrderedDict(sorted(int_inputs.items()))
+    real_inputs = OrderedDict(sorted(real_inputs.items()))
+    inputs = int_inputs.copy()
+    inputs.update(real_inputs)
+
+    log_measure = random_gaussian(inputs)
+    integrand = reduce(ops.add, [Variable(k, d) for k, d in real_inputs.items()])
+    reduced_vars = frozenset(real_inputs)
+
+    with monte_carlo_interpretation(particle=bint(100000)):
+        approx = Integrate(log_measure, integrand, reduced_vars)
+        assert isinstance(approx, Tensor)
+
+    exact = Integrate(log_measure, integrand, reduced_vars)
+    assert isinstance(exact, Tensor)
+    assert_close(approx, exact, atol=0.1, rtol=0.1)
+
+
+@pytest.mark.parametrize('int_inputs', [
+    {},
+    {'i': bint(2)},
+    {'i': bint(2), 'j': bint(3)},
+], ids=id_from_inputs)
+@pytest.mark.parametrize('real_inputs', [
+    {'x': reals()},
+    {'x': reals(2)},
+    {'x': reals(), 'y': reals()},
+    {'x': reals(2), 'y': reals(3)},
+    {'x': reals(4), 'y': reals(2, 3)},
+], ids=id_from_inputs)
+def test_integrate_gaussian(int_inputs, real_inputs):
+    int_inputs = OrderedDict(sorted(int_inputs.items()))
+    real_inputs = OrderedDict(sorted(real_inputs.items()))
+    inputs = int_inputs.copy()
+    inputs.update(real_inputs)
+
+    log_measure = random_gaussian(inputs)
+    integrand = random_gaussian(inputs)
+    reduced_vars = frozenset(real_inputs)
+
+    with monte_carlo_interpretation(particle=bint(10000)):
+        approx = Integrate(log_measure, integrand, reduced_vars)
+        assert isinstance(approx, Tensor)
+
+    exact = Integrate(log_measure, integrand, reduced_vars)
+    assert isinstance(exact, Tensor)
+    assert_close(approx, exact, atol=0.1, rtol=0.1)
diff --git a/test/test_samplers.py b/test/test_samplers.py
index 9764ca6..3af0ad3 100644
--- a/test/test_samplers.py
+++ b/test/test_samplers.py
@@ -9,6 +9,7 @@ from torch.autograd import grad
 
 import funsor.ops as ops
 from funsor.domains import bint, reals
+from funsor.integrate import Integrate
 from funsor.joint import Joint
 from funsor.terms import Variable
 from funsor.testing import assert_close, id_from_inputs, random_gaussian, random_tensor
@@ -152,7 +153,7 @@ def test_tensor_distribution(event_inputs, batch_inputs, test_grad):
     p = random_tensor(be_inputs)
     p.data.requires_grad_(test_grad)
 
-    q = p.sample(sampled_vars, sample_inputs) - ops.log(num_samples)
+    q = p.sample(sampled_vars, sample_inputs)
     mq = materialize(q).reduce(ops.logaddexp, 'n')
     mq = mq.align(tuple(p.inputs))
     assert_close(mq, p, atol=0.1, rtol=None)
@@ -166,42 +167,38 @@ def test_tensor_distribution(event_inputs, batch_inputs, test_grad):
         assert_close(actual, expected, atol=0.1, rtol=None)
 
 
-# This is a stub for a future PR.
-def Integrate(log_measure, integrand, reduced_vars):
-    pytest.xfail(reason='Integrate is not implemented')
-
-
 @pytest.mark.parametrize('batch_inputs', [
     (),
-    (('b', bint(4)),),
-    (('b', bint(4)), ('c', bint(5))),
+    (('b', bint(3)),),
+    (('b', bint(3)), ('c', bint(4))),
 ], ids=id_from_inputs)
 @pytest.mark.parametrize('event_inputs', [
     (('e', reals()),),
     (('e', reals()), ('f', reals(2))),
 ], ids=id_from_inputs)
 def test_gaussian_distribution(event_inputs, batch_inputs):
-    num_samples = 10000
-    sample_inputs = OrderedDict(n=bint(num_samples))
+    num_samples = 100000
+    sample_inputs = OrderedDict(particle=bint(num_samples))
     be_inputs = OrderedDict(batch_inputs + event_inputs)
     batch_inputs = OrderedDict(batch_inputs)
     event_inputs = OrderedDict(event_inputs)
     sampled_vars = frozenset(event_inputs)
     p = random_gaussian(be_inputs)
 
-    q = p.sample(sampled_vars, sample_inputs) - ops.log(num_samples)
+    q = p.sample(sampled_vars, sample_inputs)
     p_vars = sampled_vars
-    q_vars = sampled_vars | frozenset(['n'])
+    q_vars = sampled_vars | frozenset(['particle'])
     # Check zeroth moment.
     assert_close(q.reduce(ops.logaddexp, q_vars),
-                 p.reduce(ops.logaddexp, p_vars), atol=1e-6, rtol=None)
+                 p.reduce(ops.logaddexp, p_vars), atol=1e-6)
     for k1, d1 in event_inputs.items():
         x = Variable(k1, d1)
         # Check first moments.
         assert_close(Integrate(q, x, q_vars),
-                     Integrate(p, x, p_vars), atol=1e-2, rtol=None)
-        for k2, d2 in event_inputs.item():
+                     Integrate(p, x, p_vars), atol=0.5, rtol=0.2)
+        for k2, d2 in event_inputs.items():
             y = Variable(k2, d2)
             # Check second moments.
+            continue  # FIXME: Quadratic integration is not supported:
             assert_close(Integrate(q, x * y, q_vars),
-                         Integrate(p, x * y, p_vars), atol=1e-2, rtol=None)
+                         Integrate(p, x * y, p_vars), atol=1e-2)
diff --git a/test/test_torch.py b/test/test_torch.py
index dade2ae..ac59196 100644
--- a/test/test_torch.py
+++ b/test/test_torch.py
@@ -8,7 +8,7 @@ import torch
 
 import funsor
 import funsor.ops as ops
-from funsor.domains import Domain, bint, reals
+from funsor.domains import Domain, bint, find_domain, reals
 from funsor.terms import Number, Variable
 from funsor.testing import assert_close, assert_equiv, check_funsor, random_tensor
 from funsor.torch import REDUCE_OP_TO_TORCH, Tensor, align_tensors, torch_einsum
@@ -260,6 +260,26 @@ def test_binary_funsor_funsor(symbol, dims1, dims2):
     check_funsor(actual, inputs, Domain((), dtype), expected_data)
 
 
+@pytest.mark.parametrize('output_shape2', [(), (2,), (3, 2)], ids=str)
+@pytest.mark.parametrize('output_shape1', [(), (2,), (3, 2)], ids=str)
+@pytest.mark.parametrize('inputs2', [(), ('a',), ('b', 'a'), ('b', 'c', 'a')], ids=str)
+@pytest.mark.parametrize('inputs1', [(), ('a',), ('a', 'b'), ('b', 'a', 'c')], ids=str)
+def test_binary_broadcast(inputs1, inputs2, output_shape1, output_shape2):
+    sizes = {'a': 4, 'b': 5, 'c': 6}
+    inputs1 = OrderedDict((k, bint(sizes[k])) for k in inputs1)
+    inputs2 = OrderedDict((k, bint(sizes[k])) for k in inputs2)
+    x1 = random_tensor(inputs1, reals(*output_shape1))
+    x2 = random_tensor(inputs1, reals(*output_shape2))
+
+    actual = x1 + x2
+    assert actual.output == find_domain(ops.add, x1.output, x2.output)
+
+    block = {'a': 1, 'b': 2, 'c': 3}
+    actual_block = actual(**block)
+    expected_block = Tensor(x1(**block).data + x2(**block).data)
+    assert_close(actual_block, expected_block)
+
+
 @pytest.mark.parametrize('scalar', [0.5])
 @pytest.mark.parametrize('dims', [(), ('a',), ('a', 'b'), ('b', 'a', 'c')])
 @pytest.mark.parametrize('symbol', BINARY_OPS)