Repository: pyro-ppl/funsor

Issue:
Implement a Delta distribution and generalized Delta funsor
It appears we may be able to represent joint samples as Delta distributions. This will require two new components:
- [x] #49 a `dist.Delta` distribution wrapping a univariate PyTorch-style distribution, and 
- [ ] #65 a general `Delta` funsor similar to `Tensor` and `Gaussian`.

Discussion:
fritzo: @eb8680 I think this is ready now. I got confused about `log_density`, but I believe we need to keep track of this in the `Delta` funsors because it scales `eager_subs` but not `eager_reduce`.

Patch:
diff --git a/funsor/__init__.py b/funsor/__init__.py
index 11a1812..64b0f35 100644
--- a/funsor/__init__.py
+++ b/funsor/__init__.py
@@ -5,8 +5,8 @@ from funsor.interpreter import reinterpret
 from funsor.terms import Funsor, Number, Variable, of_shape, to_funsor
 from funsor.torch import Function, Tensor, arange, function, torch_einsum
 
-from . import adjoint, contract, distributions, domains, einsum, gaussian, \
-    handlers, interpreter, minipyro, ops, terms, torch
+from . import (adjoint, contract, delta, distributions, domains, einsum, gaussian, handlers, interpreter, minipyro, ops,
+               terms, torch)
 
 __all__ = [
     'Domain',
@@ -20,6 +20,7 @@ __all__ = [
     'backward',
     'bint',
     'contract',
+    'delta',
     'distributions',
     'domains',
     'einsum',
diff --git a/funsor/delta.py b/funsor/delta.py
new file mode 100644
index 0000000..5102d38
--- /dev/null
+++ b/funsor/delta.py
@@ -0,0 +1,109 @@
+from __future__ import absolute_import, division, print_function
+
+from collections import OrderedDict
+
+from six import add_metaclass
+
+import funsor.ops as ops
+from funsor.domains import reals
+from funsor.ops import Op
+from funsor.terms import Align, Binary, Funsor, FunsorMeta, Number, Variable, eager, to_funsor
+from funsor.torch import Tensor
+
+
+class DeltaMeta(FunsorMeta):
+    """
+    Wrapper to fill in defaults.
+    """
+    def __call__(cls, name, point, log_density=0):
+        point = to_funsor(point)
+        log_density = to_funsor(log_density)
+        return super(DeltaMeta, cls).__call__(name, point, log_density)
+
+
+@add_metaclass(DeltaMeta)
+class Delta(Funsor):
+    """
+    Normalized delta distribution binding a single variable.
+
+    :param str name: Name of the bound variable.
+    :param Funsor point: Value of the bound variable.
+    :param Funsor log_density: Optional log density to be added when evaluating
+        at a point. This is needed to make :class:`Delta` closed under
+        differentiable substitution.
+    """
+    def __init__(self, name, point, log_density=0):
+        assert isinstance(name, str)
+        assert isinstance(point, Funsor)
+        assert isinstance(log_density, Funsor)
+        assert log_density.output == reals()
+        inputs = OrderedDict([(name, point.output)])
+        inputs.update(point.inputs)
+        inputs.update(log_density.inputs)
+        output = reals()
+        super(Delta, self).__init__(inputs, output)
+        self.name = name
+        self.point = point
+        self.log_density = log_density
+
+    def eager_subs(self, subs):
+        value = None
+        index_part = []
+        for k, v in subs:
+            if k in self.inputs:
+                if k == self.name:
+                    value = v
+                else:
+                    assert self.name not in v.inputs
+                    index_part.append((k, v))
+        index_part = tuple(index_part)
+        if value is None and not index_part:
+            return self
+
+        name = self.name
+        point = self.point.eager_subs(index_part)
+        log_density = self.log_density.eager_subs(index_part)
+        if value is not None:
+            if isinstance(value, Variable):
+                name = value.name
+            elif isinstance(value, (Number, Tensor)) and isinstance(point, (Number, Tensor)):
+                return (value == point).all().log() + log_density
+            else:
+                # TODO Compute a jacobian, update log_prob, and emit another Delta.
+                raise ValueError('Cannot substitute a {} into a Delta'
+                                 .format(type(value).__name__))
+        return Delta(name, point, log_density)
+
+    def eager_reduce(self, op, reduced_vars):
+        if op is ops.logaddexp:
+            if self.name in reduced_vars:
+                return Number(0)  # Deltas are normalized.
+
+        # TODO Implement ops.add to simulate .to_event().
+
+        return None  # defer to default implementation
+
+
+@eager.register(Binary, Op, Delta, (Funsor, Delta, Align))
+def eager_binary(op, lhs, rhs):
+    if op is ops.add or op is ops.sub:
+        if lhs.name in rhs.inputs:
+            rhs = rhs(**{lhs.name: lhs.point})
+            return op(lhs, rhs)
+
+    return None  # defer to default implementation
+
+
+@eager.register(Binary, Op, (Funsor, Align), Delta)
+def eager_binary(op, lhs, rhs):
+    if op is ops.add:
+        if rhs.name in lhs.inputs:
+            lhs = lhs(**{rhs.name: rhs.point})
+            return op(lhs, rhs)
+
+    return None  # defer to default implementation
+
+
+__all__ = [
+    'Delta',
+]
diff --git a/funsor/distributions.py b/funsor/distributions.py
index c958e01..26ad08c 100644
--- a/funsor/distributions.py
+++ b/funsor/distributions.py
@@ -7,6 +7,7 @@ import pyro.distributions as dist
 import torch
 from six import add_metaclass
 
+import funsor.delta
 import funsor.ops as ops
 from funsor.domains import bint, reals
 from funsor.gaussian import Gaussian
@@ -150,6 +151,19 @@ def eager_delta(v, log_density, value):
     return Tensor(data, inputs)
 
 
+@eager.register(Delta, Funsor, Funsor, Variable)
+@eager.register(Delta, Variable, Funsor, Variable)
+def eager_delta(v, log_density, value):
+    assert v.output == value.output
+    return funsor.delta.Delta(value.name, v, log_density)
+
+
+@eager.register(Delta, Variable, Funsor, Funsor)
+def eager_delta(v, log_density, value):
+    assert v.output == value.output
+    return funsor.delta.Delta(v.name, value, log_density)
+
+
 class Normal(Distribution):
     dist_class = dist.Normal
 
diff --git a/funsor/ops.py b/funsor/ops.py
index 54ed23b..2f23cf7 100644
--- a/funsor/ops.py
+++ b/funsor/ops.py
@@ -67,7 +67,14 @@ def exp(x):
 
 @Op
 def log(x):
-    return np.log(x) if isinstance(x, Number) else x.log()
+    if isinstance(x, Number):
+        return np.log(x)
+    elif isinstance(x, torch.Tensor):
+        if x.dtype in (torch.uint8, torch.long):
+            x = x.float()
+        return x.log()
+    else:
+        return x.log()
 
 
 @Op
diff --git a/funsor/terms.py b/funsor/terms.py
index 76f57a3..cc8493a 100644
--- a/funsor/terms.py
+++ b/funsor/terms.py
@@ -137,6 +137,9 @@ class Funsor(object):
     def __hash__(self):
         return id(self)
 
+    def __repr__(self):
+        return '{}({})'.format(type(self).__name__, ', '.join(map(repr, self._ast_args)))
+
     def __call__(self, *args, **kwargs):
         """
         Partially evaluates this funsor by substituting dimensions.


Test patch:
diff --git a/test/test_delta.py b/test/test_delta.py
new file mode 100644
index 0000000..abe4f9c
--- /dev/null
+++ b/test/test_delta.py
@@ -0,0 +1,51 @@
+from __future__ import absolute_import, division, print_function
+
+import pytest
+import torch
+
+import funsor.ops as ops
+from funsor.delta import Delta
+from funsor.domains import reals
+from funsor.terms import Number, Variable
+from funsor.testing import check_funsor
+from funsor.torch import Tensor
+
+
+def test_eager_subs_variable():
+    v = Variable('v', reals(3))
+    point = Tensor(torch.randn(3))
+    d = Delta('foo', v)
+    assert d(v=point) is Delta('foo', point)
+
+
+@pytest.mark.parametrize('log_density', [0, 1.234])
+def test_eager_subs_ground(log_density):
+    point1 = Tensor(torch.randn(3))
+    point2 = Tensor(torch.randn(3))
+    d = Delta('foo', point1, log_density)
+    check_funsor(d(foo=point1), {}, reals(), torch.tensor(float(log_density)))
+    check_funsor(d(foo=point2), {}, reals(), torch.tensor(float('-inf')))
+
+
+def test_add_delta_funsor():
+    x = Variable('x', reals(3))
+    y = Variable('y', reals(3))
+    d = Delta('x', y)
+
+    expr = -(1 + x ** 2).log()
+    assert d + expr is d + expr(x=y)
+    assert expr + d is expr(x=y) + d
+
+
+def test_reduce():
+    point = Tensor(torch.randn(3))
+    d = Delta('foo', point)
+    assert d.reduce(ops.logaddexp, frozenset(['foo'])) is Number(0)
+
+
+@pytest.mark.parametrize('log_density', [0, 1.234])
+def test_reduce(log_density):
+    point = Tensor(torch.randn(3))
+    d = Delta('foo', point, log_density)
+    # Note that log_density affects ground substitution but does not affect reduction.
+    assert d.reduce(ops.logaddexp, frozenset(['foo'])) is Number(0)
diff --git a/test/test_distributions.py b/test/test_distributions.py
index df86083..542940b 100644
--- a/test/test_distributions.py
+++ b/test/test_distributions.py
@@ -8,6 +8,7 @@ import torch
 
 import funsor
 import funsor.distributions as dist
+from funsor.delta import Delta
 from funsor.domains import bint, reals
 from funsor.gaussian import Gaussian
 from funsor.terms import Variable
@@ -80,6 +81,14 @@ def test_delta_density(batch_shape, event_shape):
         assert_close(actual, expected)
 
 
+def test_delta_delta():
+    v = Variable('v', reals(2))
+    point = Tensor(torch.randn(2))
+    log_density = Tensor(torch.tensor(0.5))
+    d = dist.Delta(point, log_density, v)
+    assert d is Delta('v', point, log_density)
+
+
 def test_mvn_defaults():
     loc = Variable('loc', reals())
     scale = Variable('scale', reals())