Repository: pyro-ppl/funsor

Issue:
Support torch.nn.Module with multiple outputs

Discussion:

Patch:
diff --git a/examples/ss_vae_delayed.py b/examples/ss_vae_delayed.py
index 6548979..4671887 100644
--- a/examples/ss_vae_delayed.py
+++ b/examples/ss_vae_delayed.py
@@ -23,9 +23,9 @@ class SalientEncoder(nn.Module):
     pass  # TODO
 
 
-decoder = funsor.function((), (), ())(Decoder())
-nuisance_encoder = funsor.function((), ('loc_scale',))(NuisanceEncoder())
-salient_encoder = funsor.function((), (), ())(SalientEncoder())
+decoder = funsor.torch.function((), (), ())(Decoder())
+nuisance_encoder = funsor.torch.function((), ('loc_scale',))(NuisanceEncoder())
+salient_encoder = funsor.torch.function((), (), ())(SalientEncoder())
 
 
 def model(image=None):
diff --git a/funsor/__init__.py b/funsor/__init__.py
index a65ec82..be891ff 100644
--- a/funsor/__init__.py
+++ b/funsor/__init__.py
@@ -3,14 +3,13 @@ from __future__ import absolute_import, division, print_function
 from funsor.domains import Domain, bint, find_domain, reals
 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 funsor.torch import Tensor, arange, torch_einsum
 
 from . import (adjoint, delta, distributions, domains, einsum, gaussian, handlers, interpreter, joint,
                minipyro, ops, sum_product, terms, torch)
 
 __all__ = [
     'Domain',
-    'Function',
     'Funsor',
     'Number',
     'Tensor',
@@ -24,7 +23,6 @@ __all__ = [
     'domains',
     'einsum',
     'find_domain',
-    'function',
     'gaussian',
     'handlers',
     'interpreter',
diff --git a/funsor/torch.py b/funsor/torch.py
index 7e71503..7e8439b 100644
--- a/funsor/torch.py
+++ b/funsor/torch.py
@@ -93,8 +93,9 @@ class Tensor(Funsor):
     def __init__(self, data, inputs=None, dtype="real"):
         assert isinstance(data, torch.Tensor)
         assert isinstance(inputs, tuple)
-        assert all(isinstance(d.dtype, integer_types) for k, d in inputs)
         assert len(inputs) <= data.dim()
+        for (k, d), size in zip(inputs, data.shape):
+            assert d.dtype == size
         inputs = OrderedDict(inputs)
         output = Domain(data.shape[len(inputs):], dtype)
         super(Tensor, self).__init__(inputs, output)
@@ -375,18 +376,23 @@ def materialize(x):
     return x.eager_subs(subs)
 
 
+class LazyTuple(tuple):
+    def __call__(self, *args, **kwargs):
+        return LazyTuple(x(*args, **kwargs) for x in self)
+
+
 class Function(Funsor):
     r"""
     Funsor wrapped by a PyTorch function.
 
-    Functions are support broadcasting and can be eagerly evaluated on funsors
-    with free variables of int type (i.e. batch dimensions).
+    Functions are assumed to support broadcasting and can be eagerly evaluated
+    on funsors with free variables of int type (i.e. batch dimensions).
 
-    :class:`Function`s are often created via the :func:`function` decorator.
+    :class:`Function`s are usually created via the :func:`function` decorator.
 
     :param callable fn: A PyTorch function to wrap.
     :param funsor.domains.Domain output: An output domain.
-    :param Funsor \*args: Funsor arguments.
+    :param Funsor args: Funsor arguments.
     """
     def __init__(self, fn, output, args):
         assert callable(fn)
@@ -400,12 +406,12 @@ class Function(Funsor):
         self.args = args
 
     def __repr__(self):
-        return 'Function({})'.format(', '.join(
-            [type(self).__name__, repr(self.output)] + list(map(repr, self.args))))
+        return '{}({}, {}, {})'.format(type(self).__name__, self.fn.__name__,
+                                       repr(self.output), repr(self.args))
 
     def __str__(self):
-        return 'Function({})'.format(', '.join(
-            [type(self).__name__, str(self.output)] + list(map(str, self.args))))
+        return '{}({}, {}, {})'.format(type(self).__name__, self.fn.__name__,
+                                       str(self.output), str(self.args))
 
     def eager_subs(self, subs):
         if not any(k in self.inputs for k, v in subs):
@@ -425,11 +431,54 @@ def eager_function(fn, output, args):
     return result
 
 
+def _select(fn, i, *args):
+    result = fn(*args)
+    assert isinstance(result, tuple)
+    return result[i]
+
+
+def _nested_function(fn, args, output):
+    if isinstance(output, Domain):
+        return Function(fn, output, args)
+    elif isinstance(output, tuple):
+        result = []
+        for i, output_i in enumerate(output):
+            fn_i = functools.partial(_select, fn, i)
+            fn_i.__name__ = "{}_{}".format(fn_i, i)
+            result.append(_nested_function(fn_i, args, output_i))
+        return LazyTuple(result)
+    raise TypeError("Invalid output: {}".format(output))
+
+
+class _Memoized(object):
+    def __init__(self, fn):
+        self.fn = fn
+        self._cache = None
+
+    def __call__(self, *args):
+        if self._cache is not None:
+            old_args, old_result = self._cache
+            if all(x is y for x, y in zip(args, old_args)):
+                return old_result
+        result = self.fn(*args)
+        self._cache = args, result
+        return result
+
+    @property
+    def __name__(self):
+        return self.fn.__name__
+
+
 def _function(inputs, output, fn):
     names = getargspec(fn)[0]
     args = tuple(Variable(name, domain) for (name, domain) in zip(names, inputs))
     assert len(args) == len(inputs)
-    return Function(fn, output, args)
+    if not isinstance(output, Domain):
+        assert isinstance(output, tuple)
+        # Memoize multiple-output functions so that invocations can be shared among
+        # all outputs. This is not foolproof, but does work in simple situations.
+        fn = _Memoized(fn)
+    return _nested_function(fn, args, output)
 
 
 def function(*signature):
@@ -438,21 +487,29 @@ def function(*signature):
 
     Example::
 
-        @funsor.function(reals(3,4), reals(4,5), reals(3,5))
+        @funsor.torch.function(reals(3,4), reals(4,5), reals(3,5))
         def matmul(x, y):
             return torch.matmul(x, y)
 
-        @funsor.function(reals(10), reals(10, 10), reals())
+        @funsor.torch.function(reals(10), reals(10, 10), reals())
         def mvn_log_prob(loc, scale_tril, x):
             d = torch.distributions.MultivariateNormal(loc, scale_tril)
             return d.log_prob(x)
 
+    To support functions that output nested tuples of tensors, specify a nested
+    tuple of output types, for example:
+
+        @funsor.torch.function(reals(8), (reals(), bint(8)))
+        def max_and_argmax(x):
+            return torch.max(x, dim=-1)
+
     :param \*signature: A sequence if input domains followed by a final output
         domain.
     """
     assert signature
-    assert all(isinstance(d, Domain) for d in signature)
     inputs, output = signature[:-1], signature[-1]
+    assert all(isinstance(d, Domain) for d in inputs)
+    assert isinstance(output, (Domain, tuple))
     return functools.partial(_function, inputs, output)
 
 
@@ -479,7 +536,9 @@ def torch_einsum(equation, *operands):
     return Function(fn, output, operands)
 
 
+################################################################################
 # Register Ops
+################################################################################
 
 @ops.abs.register(torch.Tensor)
 def _abs(x):
@@ -585,13 +644,13 @@ REDUCE_OP_TO_TORCH = {
 
 
 __all__ = [
-    'REDUCE_OP_TO_TORCH',
     'Function',
+    'REDUCE_OP_TO_TORCH',
     'Tensor',
     'align_tensor',
     'align_tensors',
     'arange',
-    'torch_einsum',
     'function',
     'materialize',
+    'torch_einsum',
 ]


Test patch:
diff --git a/test/test_distributions.py b/test/test_distributions.py
index 8f9de6a..d72ec1a 100644
--- a/test/test_distributions.py
+++ b/test/test_distributions.py
@@ -59,7 +59,7 @@ def test_delta_density(batch_shape, event_shape):
     batch_dims = ('i', 'j', 'k')[:len(batch_shape)]
     inputs = OrderedDict((k, bint(v)) for k, v in zip(batch_dims, batch_shape))
 
-    @funsor.function(reals(*event_shape), reals(), reals(*event_shape), reals())
+    @funsor.torch.function(reals(*event_shape), reals(), reals(*event_shape), reals())
     def delta(v, log_density, value):
         eq = (v == value)
         for _ in range(len(event_shape)):
@@ -198,7 +198,7 @@ def test_mvn_density(batch_shape):
     batch_dims = ('i', 'j', 'k')[:len(batch_shape)]
     inputs = OrderedDict((k, bint(v)) for k, v in zip(batch_dims, batch_shape))
 
-    @funsor.function(reals(3), reals(3, 3), reals(3), reals())
+    @funsor.torch.function(reals(3), reals(3, 3), reals(3), reals())
     def mvn(loc, scale_tril, value):
         return torch.distributions.MultivariateNormal(loc, scale_tril=scale_tril).log_prob(value)
 
diff --git a/test/test_torch.py b/test/test_torch.py
index 684c6a9..2bb6cb6 100644
--- a/test/test_torch.py
+++ b/test/test_torch.py
@@ -384,7 +384,7 @@ def test_all_equal(shape):
 
 def test_function_matmul():
 
-    @funsor.function(reals(3, 4), reals(4, 5), reals(3, 5))
+    @funsor.torch.function(reals(3, 4), reals(4, 5), reals(3, 5))
     def matmul(x, y):
         return torch.matmul(x, y)
 
@@ -399,15 +399,15 @@ def test_function_matmul():
 
 def test_function_lazy_matmul():
 
-    @funsor.function(reals(3, 4), reals(4, 5), reals(3, 5))
+    @funsor.torch.function(reals(3, 4), reals(4, 5), reals(3, 5))
     def matmul(x, y):
         return torch.matmul(x, y)
 
-    x_lazy = funsor.Variable('x', reals(3, 4))
+    x_lazy = Variable('x', reals(3, 4))
     y = Tensor(torch.randn(4, 5))
     actual_lazy = matmul(x_lazy, y)
     check_funsor(actual_lazy, {'x': reals(3, 4)}, reals(3, 5))
-    assert isinstance(actual_lazy, funsor.Function)
+    assert isinstance(actual_lazy, funsor.torch.Function)
 
     x = Tensor(torch.randn(3, 4))
     actual = actual_lazy(x=x)
@@ -415,6 +415,45 @@ def test_function_lazy_matmul():
     check_funsor(actual, {}, reals(3, 5), expected_data)
 
 
+def test_function_nested_eager():
+
+    @funsor.torch.function(reals(8), (reals(), bint(8)))
+    def max_and_argmax(x):
+        return torch.max(x, dim=-1)
+
+    inputs = OrderedDict([('i', bint(2)), ('j', bint(3))])
+    x = Tensor(torch.randn(2, 3, 8), inputs)
+    m, a = x.data.max(dim=-1)
+    expected_max = Tensor(m, inputs, 'real')
+    expected_argmax = Tensor(a, inputs, 8)
+
+    actual_max, actual_argmax = max_and_argmax(x)
+    assert_close(actual_max, expected_max)
+    assert_close(actual_argmax, expected_argmax)
+
+
+def test_function_nested_lazy():
+
+    @funsor.torch.function(reals(8), (reals(), bint(8)))
+    def max_and_argmax(x):
+        return torch.max(x, dim=-1)
+
+    x_lazy = Variable('x', reals(8))
+    lazy_max, lazy_argmax = max_and_argmax(x_lazy)
+    assert isinstance(lazy_max, funsor.torch.Function)
+    assert isinstance(lazy_argmax, funsor.torch.Function)
+    check_funsor(lazy_max, {'x': reals(8)}, reals())
+    check_funsor(lazy_argmax, {'x': reals(8)}, bint(8))
+
+    inputs = OrderedDict([('i', bint(2)), ('j', bint(3))])
+    y = Tensor(torch.randn(2, 3, 8), inputs)
+    actual_max = lazy_max(x=y)
+    actual_argmax = lazy_argmax(x=y)
+    expected_max, expected_argmax = max_and_argmax(y)
+    assert_close(actual_max, expected_max)
+    assert_close(actual_argmax, expected_argmax)
+
+
 def test_align():
     x = Tensor(torch.randn(2, 3, 4), OrderedDict([
         ('i', bint(2)),