Repository: pyro-ppl/funsor

Issue:
Nondeterministic test failure in test_einsum.py
I suspect this is due to nondeterminism somewhere, like use of a dict rather than an `OrderedDict`.
```
==== FAILURES ====
____ test_einsum[a,ab,bc,cd->] ____

equation = 'a,ab,bc,cd->'

    @pytest.mark.parametrize('equation', EINSUM_EXAMPLES + XFAIL_EINSUM_EXAMPLES)
    def test_einsum(equation):
        inputs, outputs, operands, sizes = make_example(equation)
        funsor_operands = [
            funsor.Tensor(operand, OrderedDict([(d, funsor.bint(sizes[d])) for d in inp]))
            for inp, operand in zip(inputs, operands)
        ]
        expected = torch.einsum(equation, operands)
        actual = naive_einsum(equation, *funsor_operands)
        assert expected.shape == actual.data.shape
>       assert torch.allclose(expected, actual.data)
E       AssertionError: assert False
E        +  where False = <built-in method allclose of type object at 0x11953f470>(tensor(-0.0174), tensor(-0.0174))
E        +    where <built-in method allclose of type object at 0x11953f470> = torch.allclose
E        +    and   tensor(-0.0174) = Tensor(-0.01743006706237793, OrderedDict(), 'real').data

test/test_einsum.py:76: AssertionError
==== 1 failed, 785 passed, 9 skipped, 16 xfailed in 3.28 seconds ====
make: *** [test] Error 1
```

Discussion:
fritzo: @neerajprad can you please confirm this is the proper way to fix the rng seed in PyTest? I tried to follow a much-simplified form of how we do this in Pyro.

Patch:
diff --git a/funsor/__init__.py b/funsor/__init__.py
index f70a8d5..45c7281 100644
--- a/funsor/__init__.py
+++ b/funsor/__init__.py
@@ -1,9 +1,9 @@
 from __future__ import absolute_import, division, print_function
 
-from funsor.domains import Domain, find_domain, bint, reals
+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
+from funsor.torch import Function, Tensor, arange, einsum, function
 
 from . import distributions, domains, handlers, interpreter, minipyro, ops, terms, torch
 
@@ -18,6 +18,7 @@ __all__ = [
     'backward',
     'distributions',
     'domains',
+    'einsum',
     'find_domain',
     'function',
     'handlers',
diff --git a/funsor/torch.py b/funsor/torch.py
index 459b049..846c04c 100644
--- a/funsor/torch.py
+++ b/funsor/torch.py
@@ -7,7 +7,7 @@ import torch
 from six import add_metaclass, integer_types
 
 import funsor.ops as ops
-from funsor.domains import Domain, find_domain, bint
+from funsor.domains import Domain, bint, find_domain, reals
 from funsor.six import getargspec
 from funsor.terms import Binary, Funsor, FunsorMeta, Number, Variable, eager, to_funsor
 
@@ -377,11 +377,35 @@ def function(*signature):
     return functools.partial(_function, inputs, output)
 
 
+def einsum(equation, *operands):
+    """
+    Wrapper around :func:`torch.einsum` to operate on real-valued Funsors.
+
+    Note this operates only on the ``output`` tensor. To perform sum-product
+    contractions on named dimensions, instead use ``+`` and
+    :class:`~funsor.terms.Reduce`.
+    """
+    assert isinstance(equation, str)
+    assert isinstance(operands, tuple)
+    for x in operands:
+        assert isinstance(x, Funsor)
+        assert x.dtype == 'real'
+    inputs, output = equation.split('->')
+    inputs = inputs.split(',')
+    sizes = {dim: size
+             for input_, operand in zip(inputs, operands)
+             for dim, size in zip(input_, operand.output.shape)}
+    output = reals(*(sizes[dim] for dim in output))
+    fn = functools.partial(torch.einsum, equation)
+    return Function(fn, output, operands)
+
+
 __all__ = [
     'Function',
     'Tensor',
     'align_tensors',
     'arange',
+    'einsum',
     'function',
     'materialize',
 ]


Test patch:
diff --git a/test/conftest.py b/test/conftest.py
new file mode 100644
index 0000000..511326d
--- /dev/null
+++ b/test/conftest.py
@@ -0,0 +1,7 @@
+from __future__ import absolute_import, division, print_function
+
+import pyro
+
+
+def pytest_runtest_setup(item):
+    pyro.set_rng_seed(0)
diff --git a/test/test_torch.py b/test/test_torch.py
index 17602c1..2aaa65a 100644
--- a/test/test_torch.py
+++ b/test/test_torch.py
@@ -387,3 +387,29 @@ def test_align():
         for j in range(3):
             for k in range(4):
                 assert x(i=i, j=j, k=k) == y(i=i, j=j, k=k)
+
+
+@pytest.mark.parametrize('equation', [
+    'a->a',
+    'a,a->a',
+    'a,b->',
+    'a,b->a',
+    'a,b->b',
+    'a,b->ab',
+    'a,b->ba',
+    'ab,ba->',
+    'ab,ba->a',
+    'ab,ba->b',
+    'ab,ba->ab',
+    'ab,ba->ba',
+    'ab,bc->ac',
+])
+def test_einsum(equation):
+    sizes = dict(a=2, b=3, c=4)
+    inputs, outputs = equation.split('->')
+    inputs = inputs.split(',')
+    tensors = [torch.randn(tuple(sizes[d] for d in dims)) for dims in inputs]
+    funsors = [Tensor(x) for x in tensors]
+    expected = Tensor(torch.einsum(equation, *tensors))
+    actual = funsor.einsum(equation, *funsors)
+    assert_close(actual, expected)