Refactor Affine into Contraction

funsor/__init__.py

@@ -8,7 +8,6 @@
 
 from . import (
     adjoint,
-    affine,
     cnf,
     contract,
     delta,
@@ -38,7 +37,6 @@
     'Tensor',
     'Variable',
     'adjoint',
-    'affine',
     'arange',
     'backward',
     'bint',

funsor/cnf.py

@@ -13,7 +13,8 @@
 from funsor.gaussian import Gaussian, sym_inverse
 from funsor.interpreter import recursion_reinterpret
 from funsor.ops import AssociativeOp, DISTRIBUTIVE_OPS
-from funsor.terms import Binary, Funsor, Number, Reduce, Subs, Unary, eager, moment_matching, normalize, reflect
+from funsor.terms import Binary, Funsor, Number, Reduce, Subs, Unary, Variable, \
+    eager, moment_matching, normalize, reflect
 from funsor.torch import Tensor
 
 
@@ -65,6 +66,20 @@ def __init__(self, red_op, bin_op, reduced_vars, terms):
         self.bin_op = bin_op
         self.terms = terms
         self.reduced_vars = reduced_vars
+        self.is_affine = self._is_affine()
+
+    def _is_affine(self):
+        for t in self.terms:
+            if not isinstance(t, (Number, Tensor, Variable, Contraction)):
+                return False
+            if isinstance(t, Contraction):
+                if not (self.bin_op, t.bin_op) in DISTRIBUTIVE_OPS and t.is_affine:
+                    return False
+
+        if self.bin_op is ops.add and self.red_op is not anyop:
+            return sum(1 for k, v in self.inputs.items() if v.dtype == 'real') == \
+                sum(sum(1 for k, v in t.inputs.items() if v.dtype == 'real') for t in self.terms)
+        return True
 
 
 @recursion_reinterpret.register(Contraction)
@@ -236,15 +251,29 @@ def binary_subtract(op, lhs, rhs):
     return lhs + -rhs
 
 
+@normalize.register(Binary, ops.DivOp, Funsor, Funsor)
+def binary_divide(op, lhs, rhs):
+    return lhs * Unary(ops.reciprocal, rhs)
+
+
 @normalize.register(Unary, ops.NegOp, Contraction)
 def unary_contract(op, arg):
     if arg.bin_op is ops.add and arg.red_op is anyop:
-        return Contraction(arg.red_op, arg.bin_op, arg.reduced_vars, *(-t for t in arg.terms))
-    raise NotImplementedError("TODO")
+        return Contraction(arg.red_op, arg.bin_op, arg.reduced_vars, *(op(t) for t in arg.terms))
+    if arg.bin_op is ops.mul:
+        return arg * Number(-1.)
+    return None
+
+
+@normalize.register(Unary, ops.NegOp, Variable)
+def unary_neg_variable(op, arg):
+    return arg * Number(-1.)
 
 
 @normalize.register(Unary, ops.ReciprocalOp, Contraction)
 def unary_contract(op, arg):
+    if arg.bin_op is ops.mul and arg.red_op is anyop:
+        return Contraction(arg.red_op, arg.bin_op, arg.reduced_vars, *(op(t) for t in arg.terms))
     raise NotImplementedError("TODO")
 
 

funsor/distributions.py

@@ -10,7 +10,7 @@
 
 import funsor.delta
 import funsor.ops as ops
-from funsor.affine import Affine
+from funsor.cnf import Contraction
 from funsor.domains import bint, reals
 from funsor.gaussian import BlockMatrix, BlockVector, Gaussian
 from funsor.interpreter import interpretation
@@ -392,16 +392,28 @@ def eager_normal(loc, scale, value):
     return Normal(loc, scale, 'value')(value=value)
 
 
-@eager.register(Normal, (Variable, Affine), Tensor, (Variable, Affine))
-@eager.register(Normal, (Variable, Affine), Tensor, Tensor)
-@eager.register(Normal, Tensor, Tensor, (Variable, Affine))
+@eager.register(Normal, (Variable, Contraction), Tensor, (Variable, Contraction))
+@eager.register(Normal, (Variable, Contraction), Tensor, Tensor)
+@eager.register(Normal, Tensor, Tensor, (Variable, Contraction))
 def eager_normal(loc, scale, value):
     affine = (loc - value) / scale
-    assert isinstance(affine, Affine)
+    if not affine.is_affine:
+        return None
+
     real_inputs = OrderedDict((k, v) for k, v in affine.inputs.items() if v.dtype == 'real')
     assert not any(v.shape for v in real_inputs.values())
 
-    tensors = [affine.const] + [c for v, c in affine.coeffs.items()]
+    const, coeffs = to_funsor(torch.tensor(0.)), OrderedDict((k, Number(0.)) for k in real_inputs)
+    for t in affine.terms:
+        if isinstance(t, (Number, Tensor)):
+            const += t
+        elif isinstance(t, Variable):
+            coeffs[t.name] += 1.
+        elif isinstance(t, Contraction):
+            v, c = t.terms if isinstance(t.terms[0], Variable) else reversed(t.terms)
+            coeffs[v.name] += c
+
+    tensors = [const] + list(coeffs.values())
     inputs, tensors = align_tensors(*tensors)
     tensors = torch.broadcast_tensors(*tensors)
     const, coeffs = tensors[0], tensors[1:]

funsor/ops.py

@@ -72,6 +72,10 @@ class NegOp(Op):
     pass
 
 
+class DivOp(Op):
+    pass
+
+
 class GetitemMeta(type):
     _cache = {}
 
@@ -112,7 +116,7 @@ def _default(self, x, y):
 ne = Op(operator.ne)
 neg = NegOp(operator.neg)
 sub = SubOp(operator.sub)
-truediv = Op(operator.truediv)
+truediv = DivOp(operator.truediv)
 
 add = AddOp(operator.add)
 and_ = AssociativeOp(operator.and_)
@@ -199,13 +203,13 @@ def logaddexp(x, y):
     return log(exp(x - shift) + exp(y - shift)) + shift
 
 
-@Op
+@SubOp
 def safesub(x, y):
     if isinstance(y, Number):
         return sub(x, y)
 
 
-@Op
+@DivOp
 def safediv(x, y):
     if isinstance(y, Number):
         return truediv(x, y)

funsor/terms.py

@@ -1218,6 +1218,11 @@ def _log1p(x):
     return Unary(ops.log1p, x)
 
 
+@ops.reciprocal.register(Funsor)
+def _reciprocal(x):
+    return Unary(ops.reciprocal, x)
+
+
 __all__ = [
     'Binary',
     'Funsor',