Repository: pyro-ppl/funsor

Issue:
Fix docs for initial release
## Tasks
- [x] #200 Fix readthedocs build
- [x] #202 ensure all modules appear in index.rst
- [x] #277 add docstrings to funsor.pyro.hmm
- [x] #280 ensure all user-facing Funsor constructors are documented (i.e. with `:param: ...`)
- [ ] #283 update README.md, especially the section on organization

## Triaged
- work around lack of support for default options (commented out in #200)

Discussion:

Patch:
diff --git a/README.md b/README.md
index 790bc45..e53a624 100644
--- a/README.md
+++ b/README.md
@@ -4,23 +4,70 @@
 
 # Funsor
 
-Functional analysis + tensors + symbolic algebra.
+Funsor is a tensor-like library for functions and distributions.
 
-This library is an experimental work in progress.
-Beware building on top of this unstable prototype.
+See
+[Functional tensors for probabilistic programming](https://openreview.net/pdf?id=HkecHuIaUS)
+for a system description.
+
+## Installing
+
+**Install using pip:**
+
+Funsor supports Python 3.6+.
+
+```sh
+pip install funsor@https://api.github.com/repos/pyro-ppl/funsor/tarball/master
+```
+
+**Install from source:**
+```sh
+git clone git@github.com:pyro-ppl/funsor.git
+cd funsor
+git checkout master
+pip install .
+```
+
+## Using funsor
+
+Funsor can be used through a number of interfaces:
+
+-   Funsors can be used directly for probabilistic computations, using PyTorch
+    optimizers in a standard training loop. Start with these examples:
+    [discrete_hmm](examples/discrete_hmm.py),
+    [eeg_slds](examples/eeg_slds.py),
+    [kalman_filter](examples/kalman_filter.py),
+    [pcfg](examples/pcfg.py),
+    [sensor](examples/sensor.py),
+    [slds](examples/slds.py), and
+    [vae](examples/slds.py).
+-   Funsors can be used to implement custom inference algorithms within Pyro,
+    using custom elbo implementations in standard
+    [pyro.infer.SVI](http://docs.pyro.ai/en/stable/inference_algos.html#pyro.infer.svi.SVI)
+    training. See these examples:
+    [mixed_hmm](examples/mixed_hmm/model.py) and
+    [bart forecasting](https://github.com/pyro-ppl/sandbox/blob/master/2019-08-time-series/bart/forecast.py).
+-   [funsor.pyro](https://funsor.readthedocs.io/en/latest/pyro.html) provides a
+    number of Pyro-compatible (and PyTorch-compatible) distribution classes
+    that use funsors under the hood, as well
+    [utilities](https://funsor.readthedocs.io/en/latest/pyro.html#module-funsor.pyro.convert)
+    to convert between funsors and distributions.
+-   [funsor.minipyro](https://funsor.readthedocs.io/en/latest/minipyro.html)
+    provides a limited alternate backend for the Pyro probabilistic programming
+    language, and can perform some ELBO computations exactly.
 
 ## Design
 
-See [design doc](https://docs.google.com/document/d/1NVlfQnNQ0Aebg8vfIGcJKsnSqAhB4bbClQrb5dwm2OM).
+See [design doc](https://docs.google.com/document/d/1NVlfQnNQ0Aebg8vfIGcJKsnSqAhB4bbClQrb5dwm2OM). 
 
 The goal of this library is to generalize [Pyro](http://pyro.ai)'s delayed
 inference algorithms from discrete to continuous variables, and to create
 machinery to enable partially delayed sampling compatible with universality. To
 achieve this goal this library makes three orthogonal design choices:
 
-1.  Functions are first class objects. Funsors generalize the tensor interface
+1.  Open terms are objects. Funsors generalize the tensor interface
     to also cover arbitrary functions of multiple variables ("inputs"), where
-    variables may be integers, real numbers or themselves tensors. Function
+    variables may be integers, real numbers, or real tensors. Function
     evaluation / substitution is the basic operation, generalizing tensor
     indexing.  This allows probability distributions to be first-class Funsors
     and make use of existing tensor machinery, for example we can generalize
@@ -63,19 +110,9 @@ def pyro_sample(name, dist, obs=None):
     return value
 
 # ...later during inference...
-log_prob = trace_log_prob.reduce(logaddexp)  # collapses delayed variables
-loss = -funsor.eval(log_prob)                 # performs variable elimination
+loss = -trace_log_prob.reduce(logaddexp)  # collapses delayed variables
 ```
-See [examples/minipyro.py](examples/minipyro.py) for a more complete example.
-
-## Code organization
-
-- `funsor.ops` is a collection of basic ops: unary, binary, and reductions.
-- `funsor.terms` contains AST classes for symbolic algebra.
-- `funsor.torch` contains wrappers around PyTorch `Tensor`s and functions.
-- `funsor.distributions` contains standard probability distributions.
-- `funsor.interpreter` implements different evaluation strategies.
-- `funsor.minipyro` a small Funsor-compatible implementation of Pyro.
+See [funsor/minipyro.py](funsor/minipyro.py) for complete implementation.
 
 ## Related projects
 
diff --git a/docs/source/index.rst b/docs/source/index.rst
index 05115e8..93fee12 100644
--- a/docs/source/index.rst
+++ b/docs/source/index.rst
@@ -3,8 +3,8 @@
    You can adapt this file completely to your liking, but it should at least
    contain the root `toctree` directive.
 
-Funsor is a tensor-like library for functional analysis
-=======================================================
+Funsor is a tensor-like library for functions and distributions
+===============================================================
 
 .. toctree::
    :glob:
diff --git a/funsor/affine.py b/funsor/affine.py
index f26ab74..d9508a1 100644
--- a/funsor/affine.py
+++ b/funsor/affine.py
@@ -1,16 +1,102 @@
 from collections import OrderedDict
+from functools import reduce, singledispatch
 
 import opt_einsum
 import torch
 
-from funsor.interpreter import gensym
-from funsor.terms import Lambda, Variable, bint
-from funsor.torch import Tensor
+from funsor.cnf import Contraction
+from funsor.interpreter import gensym, interpretation
+from funsor.terms import Binary, Funsor, Lambda, Reduce, Unary, Variable, bint, reflect
+from funsor.torch import Einsum, Tensor
+
+from . import ops
 
 
-# FIXME change this to a sound but incomplete test using pattern matching.
 def is_affine(fn):
-    return True
+    """
+    A sound but incomplete test to determine whether a funsor is affine with
+    respect to all of its real inputs.
+
+    :param Funsor fn: A funsor.
+    :rtype: bool
+    """
+    assert isinstance(fn, Funsor)
+    return _affine_inputs(fn) == _real_inputs(fn)
+
+
+def _real_inputs(fn):
+    return frozenset(k for k, d in fn.inputs.items() if d.dtype == "real")
+
+
+@singledispatch
+def _affine_inputs(fn):
+    """
+    Returns a [sound sub]set of real inputs of ``fn``
+    wrt which ``fn`` is known to be affine.
+    """
+    return frozenset()
+
+
+@_affine_inputs.register(Variable)
+def _(fn):
+    return _real_inputs(fn)
+
+
+@_affine_inputs.register(Unary)
+def _(fn):
+    if fn.op in (ops.neg, ops.add) or isinstance(fn.op, ops.ReshapeOp):
+        return _affine_inputs(fn.arg)
+    return frozenset()
+
+
+@_affine_inputs.register(Binary)
+def _(fn):
+    if fn.op in (ops.add, ops.sub):
+        return _affine_inputs(fn.lhs) | _affine_inputs(fn.rhs)
+    if fn.op is ops.truediv:
+        return _affine_inputs(fn.lhs) - _real_inputs(fn.rhs)
+    if isinstance(fn.op, ops.GetitemOp):
+        return _affine_inputs(fn.lhs)
+    if fn.op in (ops.mul, ops.matmul):
+        lhs_affine = _affine_inputs(fn.lhs) - _real_inputs(fn.rhs)
+        rhs_affine = _affine_inputs(fn.rhs) - _real_inputs(fn.lhs)
+        if not lhs_affine:
+            return rhs_affine
+        if not rhs_affine:
+            return lhs_affine
+        # This multilinear case introduces incompleteness, since some vars
+        # could later be reduced, making remaining vars affine.
+        return frozenset()
+    return frozenset()
+
+
+@_affine_inputs.register(Reduce)
+def _(fn):
+    return _affine_inputs(fn.arg) - fn.reduced_vars
+
+
+@_affine_inputs.register(Contraction)
+def _(fn):
+    with interpretation(reflect):
+        flat = reduce(fn.bin_op, fn.terms).reduce(fn.red_op, fn.reduced_vars)
+    return _affine_inputs(flat)
+
+
+@_affine_inputs.register(Einsum)
+def _(fn):
+    # This is simply a multiary version of the above Binary(ops.mul, ...) case.
+    results = []
+    for i, x in enumerate(fn.operands):
+        others = fn.operands[:i] + fn.operands[i+1:]
+        other_inputs = reduce(ops.or_, map(_real_inputs, others), frozenset())
+        results.append(_affine_inputs(x) - other_inputs)
+    # This multilinear case introduces incompleteness, since some vars
+    # could later be reduced, making remaining vars affine.
+    if sum(map(bool, results)) == 1:
+        for result in results:
+            if result:
+                return result
+    return frozenset()
 
 
 def extract_affine(fn):
diff --git a/funsor/cnf.py b/funsor/cnf.py
index 2ff3614..28c8dcb 100644
--- a/funsor/cnf.py
+++ b/funsor/cnf.py
@@ -58,21 +58,6 @@ class Contraction(Funsor):
         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 ((t.bin_op, self.bin_op) in DISTRIBUTIVE_OPS or (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 nullop:
-            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
 
     def unscaled_sample(self, sampled_vars, sample_inputs):
         sampled_vars = sampled_vars.intersection(self.inputs)
diff --git a/funsor/distributions.py b/funsor/distributions.py
index 416e490..d70ce67 100644
--- a/funsor/distributions.py
+++ b/funsor/distributions.py
@@ -483,7 +483,7 @@ def eager_normal(loc, scale, value):
 @eager.register(Normal, Tensor, Tensor, (Variable, Contraction))
 def eager_normal(loc, scale, value):
     affine = (loc - value) / scale
-    if not affine.is_affine:
+    if not is_affine(affine):
         return None
 
     real_inputs = OrderedDict((k, v) for k, v in affine.inputs.items() if v.dtype == 'real')
diff --git a/funsor/domains.py b/funsor/domains.py
index 2679062..cd1e337 100644
--- a/funsor/domains.py
+++ b/funsor/domains.py
@@ -90,6 +90,8 @@ def find_domain(op, *domains):
             dtype = 'real'
         elif isinstance(op, ops.ReshapeOp):
             shape = op.shape
+        elif isinstance(op, ops.AssociativeOp):
+            shape = ()
         return Domain(shape, dtype)
 
     lhs, rhs = domains


Test patch:
diff --git a/test/test_affine.py b/test/test_affine.py
index 83d0895..633dbdc 100644
--- a/test/test_affine.py
+++ b/test/test_affine.py
@@ -3,13 +3,15 @@ from collections import OrderedDict
 import pytest
 import torch
 
-from funsor.affine import extract_affine
+from funsor.affine import extract_affine, is_affine
 from funsor.cnf import Contraction
 from funsor.domains import bint, reals
-from funsor.terms import Number, Variable
-from funsor.testing import assert_close, check_funsor, random_tensor
+from funsor.terms import Number, Unary, Variable
+from funsor.testing import assert_close, check_funsor, random_gaussian, random_tensor
 from funsor.torch import Einsum, Tensor
 
+assert random_gaussian  # flake8
+
 SMOKE_TESTS = [
     ('t+x', Contraction),
     ('x+t', Contraction),
@@ -39,7 +41,7 @@ def test_smoke(expr, expected_type):
 
     result = eval(expr)
     assert isinstance(result, expected_type)
-    assert result.is_affine
+    assert is_affine(result)
 
 
 SUBS_TESTS = [
@@ -75,10 +77,12 @@ def test_affine_subs(expr, expected_type, expected_inputs):
     result = eval(expr)
     assert isinstance(result, expected_type)
     check_funsor(result, expected_inputs, expected_output)
-    assert result.is_affine
+    assert is_affine(result)
 
 
 @pytest.mark.parametrize('expr', [
+    "-Variable('x', reals())",
+    "Variable('x', reals(2)).sum()",
     "Variable('x', reals()) + 0.5",
     "Variable('x', reals(2, 3)) + Variable('y', reals(2, 3))",
     "Variable('x', reals(2)) + Variable('y', reals(2))",
@@ -95,7 +99,8 @@ def test_affine_subs(expr, expected_type, expected_inputs):
 ])
 def test_extract_affine(expr):
     x = eval(expr)
-    assert isinstance(x, (Contraction, Einsum))
+    assert is_affine(x)
+    assert isinstance(x, (Unary, Contraction, Einsum))
     real_inputs = OrderedDict((k, d) for k, d in x.inputs.items()
                               if d.dtype == 'real')
 
@@ -116,3 +121,25 @@ def test_extract_affine(expr):
                          for k, (coeff, eqn) in coeffs.items())
     assert isinstance(actual, Tensor)
     assert_close(actual, expected)
+
+
+@pytest.mark.parametrize("expr", [
+    "Variable('x', reals()).log()",
+    "Variable('x', reals()).exp()",
+    "Variable('x', reals()).sigmoid()",
+    "Variable('x', reals(2)).prod()",
+    "Variable('x', reals()) ** 2",
+    "Variable('x', reals()) ** 2",
+    "2 ** Variable('x', reals())",
+    "Variable('x', reals()) * Variable('x', reals())",
+    "Variable('x', reals()) * Variable('y', reals())",
+    "Variable('x', reals()) / Variable('y', reals())",
+    "Variable('x', reals()) / Variable('y', reals())",
+    "Variable('x', reals(2,3)) @ Variable('y', reals(3,4))",
+    "random_gaussian(OrderedDict(x=reals()))",
+    "Einsum('abcd,ac->bd',"
+    " (Variable('y', reals(2, 3, 4, 5)), Variable('x', reals(2, 4))))",
+])
+def test_not_is_affine(expr):
+    x = eval(expr)
+    assert not is_affine(x)