Refactor contract dependencies

funsor/contract.py

@@ -1,19 +1,12 @@
 from __future__ import absolute_import, division, print_function
 
+import functools
 from collections import OrderedDict
 
-import opt_einsum
-
 import funsor.ops as ops
-from funsor.distributions import Gaussian, Delta
 from funsor.optimizer import Finitary, optimize
 from funsor.sum_product import _partition
-from funsor.terms import Funsor, Number, Variable, eager
-from funsor.torch import Tensor
-
-
-# TODO handle Joint as well
-ATOMS = (Tensor, Gaussian, Delta, Number, Variable)
+from funsor.terms import Funsor, eager
 
 
 def _order_lhss(lhs, reduced_vars):
@@ -29,10 +22,13 @@ def _order_lhss(lhs, reduced_vars):
     return root_lhs, remaining_lhs
 
 
-def _simplify_contract(lhs, rhs, reduced_vars):
+def _simplify_contract(fn, lhs, rhs, reduced_vars):
     """
     Reduce free variables that do not appear explicitly in the lhs
     """
+    if not reduced_vars:
+        return lhs * rhs
+
     lhs_vars = frozenset(lhs.inputs)
     rhs_vars = frozenset(rhs.inputs)
     assert reduced_vars <= lhs_vars | rhs_vars
@@ -45,11 +41,17 @@ def _simplify_contract(lhs, rhs, reduced_vars):
         lhs = lhs.reduce(ops.add, reduced_vars - rhs_vars)
         reduced_vars = reduced_vars & rhs_vars
         progress = True
-
     if progress:
         return Contract(lhs, rhs, reduced_vars)
 
-    return None
+    return fn(lhs, rhs, reduced_vars)
+
+
+def contractor(fn):
+    """
+    Decorator for contract implementations to simplify inputs.
+    """
+    return functools.partial(_simplify_contract, fn)
 
 
 class Contract(Funsor):
@@ -77,50 +79,27 @@ def eager_subs(self, subs):
                         self.reduced_vars)
 
 
-@optimize.register(Contract, ATOMS[1:], ATOMS, frozenset)
-@optimize.register(Contract, ATOMS, ATOMS[1:], frozenset)
-@eager.register(Contract, ATOMS[1:], ATOMS, frozenset)
-@eager.register(Contract, ATOMS, ATOMS[1:], frozenset)
-def contract_ground_ground(lhs, rhs, reduced_vars):
-    result = _simplify_contract(lhs, rhs, reduced_vars)
-    if result is not None:
-        return result
-
+@eager.register(Contract, Funsor, Funsor, frozenset)
+@contractor
+def eager_contract(lhs, rhs, reduced_vars):
     return (lhs * rhs).reduce(ops.add, reduced_vars)
 
 
-@eager.register(Contract, Tensor, Tensor, frozenset)
-def eager_contract_tensor_tensor(lhs, rhs, reduced_vars):
-    result = _simplify_contract(lhs, rhs, reduced_vars)
-    if result is not None:
-        return result
-
-    out_inputs = OrderedDict([(k, d) for t in (lhs, rhs)
-                              for k, d in t.inputs.items() if k not in reduced_vars])
-
-    return Tensor(
-        opt_einsum.contract(lhs.data, list(lhs.inputs.keys()),
-                            rhs.data, list(rhs.inputs.keys()),
-                            list(out_inputs.keys()), backend="torch"),
-        out_inputs
-    )
-
+@optimize.register(Contract, Funsor, Funsor, frozenset)
+@contractor
+def optmize_contract(lhs, rhs, reduced_vars):
+    return None
 
-@optimize.register(Contract, ATOMS, Finitary, frozenset)
-def contract_ground_finitary(lhs, rhs, reduced_vars):
-    result = _simplify_contract(lhs, rhs, reduced_vars)
-    if result is not None:
-        return result
 
+@optimize.register(Contract, Funsor, Finitary, frozenset)
+@contractor
+def optimize_contract_funsor_finitary(lhs, rhs, reduced_vars):
     return Contract(rhs, lhs, reduced_vars)
 
 
-@optimize.register(Contract, Finitary, (Finitary,) + ATOMS, frozenset)
-def contract_finitary_ground(lhs, rhs, reduced_vars):
-    result = _simplify_contract(lhs, rhs, reduced_vars)
-    if result is not None:
-        return result
-
+@optimize.register(Contract, Finitary, (Finitary, Funsor), frozenset)
+@contractor
+def optimize_contract_finitary_funsor(lhs, rhs, reduced_vars):
     # exploit linearity of contraction
     if lhs.op is ops.add:
         return Finitary(
@@ -139,3 +118,9 @@ def contract_finitary_ground(lhs, rhs, reduced_vars):
                             reduced_vars & frozenset(root_lhs.inputs))
 
     return None
+
+
+__all__ = [
+    'Contract',
+    'contractor',
+]

funsor/torch.py

@@ -3,11 +3,13 @@
 import functools
 from collections import OrderedDict
 
+import opt_einsum
 import torch
 from six import add_metaclass, integer_types
 from six.moves import reduce
 
 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.ops import Op
@@ -352,6 +354,18 @@ def eager_binary_tensor_tensor(op, lhs, rhs):
     return Tensor(data, inputs, dtype)
 
 
+@eager.register(Contract, Tensor, Tensor, frozenset)
+@contractor
+def eager_contract(lhs, rhs, reduced_vars):
+    inputs = OrderedDict((k, d) for t in (lhs, rhs)
+                         for k, d in t.inputs.items() if k not in reduced_vars)
+    data = opt_einsum.contract(lhs.data, list(lhs.inputs),
+                               rhs.data, list(rhs.inputs),
+                               list(inputs), backend="torch")
+    dtype = find_domain(ops.mul, lhs.output, rhs.output).dtype
+    return Tensor(data, inputs, dtype)
+
+
 def arange(name, size):
     """
     Helper to create a named :func:`torch.arange` funsor.