Repository: pyro-ppl/pyro
Pull request: Add naive_ubersum reference implementation + tests
Addresses #1391
Fixes #1414

This clarifies the semantics of `ubersum()` by adding a simple reference implementation. The `naive_ubersum()` reference implementation is slow but handles both polynomial-time and exponential-time contractions, whereas `ubersum()` currently raises `NotImplementedError` for exponential-time contractions. It's actually very very slow; it is good for testing but wouldn't be able to serve as a fallback implementation for `ubersum()`'s NotImplemented cases.

This PR also fixes a shape bug in the `torch_log` backend (causing #1414) and adds regression tests.

## Tested

- [x] add test for agreement between `contract()` and `naive_ubersum()`
- [x] add test for agreement between `ubersum()` and `naive_ubersum()`
- [x] parametrize existing tests to also test `naive_ubersum()`
- [x] add tests for preserved batched dimensions
- [x] mark `ubersum()` discrepancies as xfail

Problem statement:
test_contract failing on PyTorch master
Some of our einsum tests (`test_ubersum_sizes`) are failing on PyTorch master:

e.g.

```python
X = torch.randn(2, 1)
Y = torch.randn(1, 3)
Z = torch.randn(3, 4)
ubersum('ab,bc,cd->a,b,c,d', X, Y, Z, batch_dims='ad')
```

throws an exception,

```
        if len(operands) == 1 and isinstance(operands[0], (list, tuple)):
            # the old interface of passing the operands as one list argument
            operands = operands[0]
>       return torch._C._VariableFunctions.einsum(equation, operands)
E       RuntimeError: size of dimension does not match previous size, operand 1, dim 0
```

The issue seems to be that PyTorch release does implicit broadcasting of operands whereas the new einsum does not. e.g.

**PyTorch Release**

```python
In [7]: torch.einsum('a,a->', [torch.tensor([1.]), torch.tensor([1., 1.])])
Out[7]: tensor(2.)
```

**PyTorch Master:**

```python
In [2]: torch.einsum('a,a->', [torch.tensor([1.]), torch.tensor([1., 1.])])
```
```
---------------------------------------------------------------------------
RuntimeError                              Traceback (most recent call last)
<ipython-input-2-65721fd81f1e> in <module>()
----> 1 torch.einsum('a,a->', [torch.tensor([1.]), torch.tensor([1., 1.])])

~/miniconda2/envs/pytorch-master/lib/python3.6/site-packages/torch/functional.py in einsum(equation, *operands)
    239         # the old interface of passing the operands as one list argument
    240         operands = operands[0]
--> 241     return torch._C._VariableFunctions.einsum(equation, operands)
    242
    243

RuntimeError: size of dimension does not match previous size, operand 1, dim 0
```

cc. @fritzo

Hints from issue discussion:

Commit message: Add naive_ubersum reference implementation and tests

Patch to review:
diff --git a/pyro/ops/contract.py b/pyro/ops/contract.py
index 141a650716..c4c6ea7547 100644
--- a/pyro/ops/contract.py
+++ b/pyro/ops/contract.py
@@ -1,11 +1,15 @@
 from __future__ import absolute_import, division, print_function
 
+import itertools
+import logging
 from abc import ABCMeta, abstractmethod
 from collections import OrderedDict, defaultdict
 
+import opt_einsum
 import torch
 from opt_einsum import shared_intermediates
 from six import add_metaclass
+from six.moves import map
 
 from pyro.distributions.util import broadcast_shape
 from pyro.ops.einsum import contract
@@ -308,11 +312,12 @@ def _contract_component(ring, tensor_tree, sum_dims):
             else:
                 parent = frozenset.union(*(t for t, d in dims_tree.items() if d & remaining_dims))
                 if parent == leaf:
-                    raise ValueError("Expected tree-structured iarange nesting, but found "
-                                     "dependencies on independent iaranges [{}]. "
-                                     "Try converting one of the iaranges to an irange (but beware "
-                                     "exponential cost in the size of that irange)"
-                                     .format(', '.join(getattr(f, 'name', str(f)) for f in leaf)))
+                    raise NotImplementedError(
+                        "Expected tree-structured iarange nesting, but found "
+                        "dependencies on independent iaranges [{}]. "
+                        "Try converting one of the iaranges to an irange (but beware "
+                        "exponential cost in the size of that irange)"
+                        .format(', '.join(getattr(f, 'name', str(f)) for f in leaf)))
                 contract_frames = leaf - parent
                 term = ring.product(term, contract_frames)
 
@@ -487,3 +492,97 @@ def ubersum(equation, *operands, **kwargs):
                 term = term.permute(*map(dims.index, output))
             results.append(term)
     return tuple(results)
+
+
+def naive_ubersum(equation, *operands, **kwargs):
+    """
+    Slow reference implementation of :func:`ubersum`.
+
+    The two implementations should agree whenever both are implemented and
+    :func:`ubersum` does not raise an error.
+    """
+    # Evaluate each output separately.
+    inputs, outputs = equation.split('->')
+    return tuple(_naive_batch_einsum(inputs + '->' + output, *operands, **kwargs)
+                 for output in outputs.split(','))
+
+
+class _DimNormalizer(object):
+    """
+    Object to map batched dims to batches of unbatched dims.
+
+    :param dict dim_to_ordinal: a mapping from contraction dim to the set of
+        batch dims over which the contraction dim is batched.
+    """
+    def __init__(self, dim_to_ordinal):
+        self._plates = {d: tuple(sorted(ordinal)) for d, ordinal in dim_to_ordinal.items()}
+        self._symbols = map(opt_einsum.get_symbol, itertools.count())
+        self._map = {}
+
+    def __call__(self, dim, indices):
+        """
+        Assigns a normalized dim to a base dim, given indices into a batch.
+
+        :param str dim: a dimension to normalize
+        :param dict indices: a mapping from batch dimension to int
+        :return: a normalied dim
+        :rtype: str
+        """
+        plate = self._plates.get(dim, ())
+        index = tuple(indices[d] for d in plate)
+        key = dim, index
+        if key in self._map:
+            return self._map[key]
+        normal_dim = next(self._symbols)
+        self._map[key] = normal_dim
+        return normal_dim
+
+
+def _naive_batch_einsum(equation, *operands, **kwargs):
+    batch_dims = set(kwargs.pop('batch_dims', ''))
+    inputs, output = equation.split('->')
+    inputs = inputs.split(',')
+    output_dims = set(output)
+    keep_dims = output_dims - batch_dims
+    sizes = {dim: size
+             for dims, term in zip(inputs, operands)
+             for dim, size in zip(dims, term.shape)}
+
+    # It is unclear what the desired behavior is in this case.
+    if batch_dims & output_dims:
+        raise NotImplementedError
+
+    # Compute batch dims for each contraction dim.
+    # By convention a sum dim's batch context is is the intersection over all
+    # batch contexts of tensors in which the sum dim appears.
+    dim_to_ordinal = {}
+    for dims in inputs:
+        dims = set(dims)
+        sum_dims = dims - batch_dims - keep_dims
+        ordinal = dims & batch_dims
+        for dim in sum_dims:
+            dim_to_ordinal[dim] = dim_to_ordinal.get(dim, ordinal) & ordinal
+
+    # Normalize by expanding batch dimensions.
+    normalize_dim = _DimNormalizer(dim_to_ordinal)
+    normal_inputs = []
+    normal_operands = []
+    normal_output = ''.join(normalize_dim(d, {})
+                            for d in output if d not in batch_dims)
+    for base_dims, term in zip(inputs, operands):
+        local_dims = [d for d in base_dims if d in batch_dims]
+        offsets = [base_dims.index(d) - len(base_dims) for d in local_dims]
+        for index in itertools.product(*(range(sizes[d]) for d in local_dims)):
+            normal_dims = ''.join(normalize_dim(d, dict(zip(local_dims, index)))
+                                  for d in base_dims if d not in batch_dims)
+            normal_term = term
+            for offset, i in zip(offsets, index):
+                normal_term = normal_term.select(offset, i)
+            normal_inputs.append(normal_dims)
+            normal_operands.append(normal_term)
+
+    # Defer to unbatched einsum.
+    normal_equation = ','.join(normal_inputs) + '->' + normal_output
+    logging.debug('normal_equation = ' + normal_equation)
+    return opt_einsum.contract(normal_equation, *normal_operands,
+                               backend='pyro.ops.einsum.torch_log')
diff --git a/tests/infer/test_enum.py b/tests/infer/test_enum.py
index baa8397d55..51e9ff5db5 100644
--- a/tests/infer/test_enum.py
+++ b/tests/infer/test_enum.py
@@ -2325,7 +2325,7 @@ def guide(data):
 
     # But promoting both to iaranges should result in an error.
     elbo = TraceEnum_ELBO(max_iarange_nesting=2)
-    with pytest.raises(ValueError, match="Expected tree-structured iarange nesting.*"):
+    with pytest.raises(NotImplementedError, match="Expected tree-structured iarange nesting.*"):
         elbo.differentiable_loss(model_iarange_iarange, guide, data)
 
 
diff --git a/tests/ops/test_contract.py b/tests/ops/test_contract.py
index 7b4609b071..b9e78a4324 100644
--- a/tests/ops/test_contract.py
+++ b/tests/ops/test_contract.py
@@ -8,7 +8,8 @@
 import torch
 
 from pyro.distributions.util import logsumexp
-from pyro.ops.contract import UnpackedLogRing, _partition_terms, contract_tensor_tree, contract_to_tensor, ubersum
+from pyro.ops.contract import (UnpackedLogRing, _partition_terms, contract_tensor_tree, contract_to_tensor,
+                               naive_ubersum, ubersum)
 from pyro.poutine.indep_messenger import CondIndepStackFrame
 from tests.common import assert_equal
 
@@ -248,6 +249,7 @@ def test_contract_to_tensor_sizes(a, b, c, d):
 
 
 UBERSUM_EXAMPLES = [
+    # non-batched examples
     ('->', ''),
     ('a->', ''),
     ('ab->', ''),
@@ -255,13 +257,24 @@ def test_contract_to_tensor_sizes(a, b, c, d):
     ('ab,bc->ca', ''),
     ('ab,bc->a,b,c', ''),
     ('ab,bc,cd->cb,da', ''),
-    ('ab,ac->,b,c,cb,a,ca,ba,ab,ac,bac', 'a'),
-    ('e,ae,be,bce,bde->,e,a,ae,b,be,bc,bce,bd,bde', 'abcd'),
+    # batched examples
+    ('a->', 'a'),
+    (',a->', 'a'),
+    (',a,ab->,b', 'a'),
+    ('ca,ab->,b,c,bc', 'a'),
+    ('ac,bc,abc->,c', 'ab'),
+    ('a,bd,abcd->,c', 'ab'),
+    ('ac,bd,abcd->,c', 'ab'),
+    (',a,b,c,ab,ac,bc,abc->,c', 'ab'),
+    ('ab,ac->,b,c,cb,bc', 'a'),
+    ('ab,ac->a,ca,ba,ab,ac,bac', 'a'),
+    (',ad,abd,acd->,d', 'abc'),
+    ('e,ae,be,bce,bde->,e', 'abcd'),
+    ('e,ae,be,bce,bde->a,ae,b,be,bc,bce,bd,bde', 'abcd'),
 ]
 
 
-@pytest.mark.parametrize('equation,batch_dims', UBERSUM_EXAMPLES)
-def test_ubersum(equation, batch_dims):
+def make_example(equation):
     symbols = sorted(set(equation) - set(',->'))
     sizes = {dim: size for dim, size in zip(symbols, range(2, 2 + len(symbols)))}
     inputs, outputs = equation.split('->')
@@ -269,9 +282,20 @@ def test_ubersum(equation, batch_dims):
     for dims in inputs.split(','):
         shape = tuple(sizes[dim] for dim in dims)
         operands.append(torch.randn(shape))
+    return operands, sizes
+
 
-    actual = ubersum(equation, *operands, batch_dims=batch_dims)
+@pytest.mark.parametrize('equation,batch_dims', UBERSUM_EXAMPLES)
+def test_ubersum_simple(equation, batch_dims):
+    operands, sizes = make_example(equation)
 
+    try:
+        actual = ubersum(equation, *operands, batch_dims=batch_dims)
+    except NotImplementedError:
+        pytest.skip()
+
+    assert isinstance(actual, tuple)
+    inputs, outputs = equation.split('->')
     outputs = outputs.split(',')
     assert len(actual) == len(outputs)
     for output, actual_part in zip(outputs, actual):
@@ -284,6 +308,21 @@ def test_ubersum(equation, batch_dims):
             assert_equal(actual_part, expected_part)
 
 
+@pytest.mark.parametrize('equation,batch_dims', UBERSUM_EXAMPLES)
+def test_ubersum_naive(equation, batch_dims):
+    operands, sizes = make_example(equation)
+
+    try:
+        actual = ubersum(equation, *operands, batch_dims=batch_dims)
+        expected = naive_ubersum(equation, *operands, batch_dims=batch_dims)
+    except NotImplementedError:
+        pytest.skip()
+
+    assert len(actual) == len(expected)
+    for a, e in zip(actual, expected):
+        assert_equal(a, e)
+
+
 @pytest.mark.parametrize('a', [2, 1])
 @pytest.mark.parametrize('b', [3, 1])
 @pytest.mark.parametrize('c', [3, 1])
@@ -364,7 +403,7 @@ def test_ubersum_collide_error():
     x = torch.randn(a, c)
     y = torch.randn(b, d)
     z = torch.randn(a, b, c, d)
-    with pytest.raises(ValueError, match='Expected tree-structured iarange nesting'):
+    with pytest.raises(NotImplementedError, match='Expected tree-structured iarange nesting'):
         ubersum('ac,bd,abcd->', x, y, z, batch_dims='ab')
 
 


Code review comments:
diff_hunk: @@ -487,3 +492,97 @@ def ubersum(equation, *operands, **kwargs):
                 term = term.permute(*map(dims.index, output))
             results.append(term)
     return tuple(results)
+
+
+def naive_ubersum(equation, *operands, **kwargs):
+    """
+    Slow reference implementation of :func:`ubersum`.
+
+    The two implementations should agree whenever both are implemented and
+    :func:`ubersum` does not raise an error.
+    """
+    # Evaluate each output separately.
+    inputs, outputs = equation.split('->')
+    return tuple(_naive_batch_einsum(inputs + '->' + output, *operands, **kwargs)
+                 for output in outputs.split(','))
+
+
+class _DimNormalizer(object):
+    """
+    Object to map batched dims to batches of unbatched dims.
+
+    :param dict dim_to_ordinal: a mapping from contraction dim to the set of
+        batch dims over which the contraction dim is batched.
+    """
+    def __init__(self, dim_to_ordinal):
+        self._plates = {d: tuple(sorted(ordinal)) for d, ordinal in dim_to_ordinal.items()}
+        self._symbols = map(opt_einsum.get_symbol, itertools.count())
+        self._map = {}
+
+    def __call__(self, dim, indices):
+        """
+        Assigns a normalized dim to a base dim, given indices into a batch.
+
+        :param str dim: a dimension to normalize
+        :param dict indices: a mapping from batch dimension to int
+        :return: a normalied dim
+        :rtype: str
+        """
+        plate = self._plates.get(dim, ())
+        index = tuple(indices[d] for d in plate)
+        key = dim, index
+        if key in self._map:
+            return self._map[key]
+        normal_dim = next(self._symbols)
+        self._map[key] = normal_dim
+        return normal_dim
+
+
+def _naive_batch_einsum(equation, *operands, **kwargs):
+    batch_dims = set(kwargs.pop('batch_dims', ''))
+    inputs, output = equation.split('->')
+    inputs = inputs.split(',')
+    output_dims = set(output)
+    keep_dims = output_dims - batch_dims
+    sizes = {dim: size
+             for dims, term in zip(inputs, operands)
+             for dim, size in zip(dims, term.shape)}
+
+    # It is unclear what the desired behavior is in this case.
+    if batch_dims & output_dims:
+        raise NotImplementedError
path: pyro/ops/contract.py
text: @user1:
It will be nice if `naive_ubersum` could handle cases that are a superset of those handled by `ubersum` (though I suppose it depends on what is the purpose of this reference implementation). One thing we could do is to first call `ubersum` and only go forward with the naive implementation if the case returns a `NotImplementedError`.

@user1:
On second thoughts, that defeats the point of having a reference implementation. Could you explain the difficulty in handling the case where some batch dimensions are not contracted?

@author:
Ok, I've narrowed down to a single sensible behavior here. I'll implement handling for this case and push to this PR...