Repository: pyro-ppl/pyro

Issue:
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 


Discussion:

Resolving patch:
diff --git a/pyro/ops/contract.py b/pyro/ops/contract.py
--- a/pyro/ops/contract.py
+++ b/pyro/ops/contract.py
@@ -1,11 +1,14 @@
 from __future__ import absolute_import, division, print_function
 
+import itertools
 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 +311,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)
 
@@ -415,7 +419,7 @@ def contract_to_tensor(tensor_tree, sum_dims, target_ordinal, ring=None, cache=N
 
 def ubersum(equation, *operands, **kwargs):
     """
-    Generalized batched einsum via tensor message passing.
+    Generalized batched sum-product algorithm via tensor message passing.
 
     This generalizes :func:`~pyro.ops.einsum.contract` in two ways:
 
@@ -423,6 +427,10 @@ def ubersum(equation, *operands, **kwargs):
     2.  Inputs and outputs can be batched along symbols given in ``batch_dims``;
         reductions along ``batch_dims`` are product reductions.
 
+    The best way to understand this function is to try the examples below,
+    which show how :func:`ubersum` calls can be implemented as multiple calls
+    to :func:`~pyro.ops.einsum.contract` (which is generally more expensive).
+
     To illustrate multiple outputs, note that the following are equivalent::
 
         z1, z2, z3 = ubersum('ab,bc->a,b,c', x, y)  # multiple outputs
@@ -434,10 +442,36 @@ def ubersum(equation, *operands, **kwargs):
 
     To illustrate batched inputs, note that the following are equivalent::
 
+        assert len(x) == 3 and len(y) == 3
         z = ubersum('c,abc,acd->bd', w, x, y, batch_dims='a')
 
-        z = w + sum(contract('bc,cd->bd', xi, yi, backend=backend)
-                    for xi, yi in zip(x, y))
+        z = contract('c,bc,bc,bc,cd,cd,cd->bd', w, *x, *y, backend=backend)
+
+    When a non-batch dimension `i` always appears with a batch dimension `a`,
+    then `i` corresponds to a distinct symbol for each slice of `a`. Thus
+    the following are equivalent::
+
+        assert len(x) == 3 and len(y) == 3
+        z = ubersum('abi,abi->', x, y, batch_dims='a')
+
+        z = contract('bi,bj,bk,bi,bj,bk->', *x, *y, backend=backend)
+
+    When such a non-batched dimension appears in the output, it must be
+    accompanied by all of its batch dimensions, e.g. the following are
+    equivalent::
+
+        assert len(x) == 3 and len(y) == 3
+        z = ubersum('abi,abi->ai', x, y, batch_dims='a')
+
+        z0 = contract('bi,bj,bk,bi,bj,bk->i', *x, *y, backend=backend)
+        z1 = contract('bi,bj,bk,bi,bj,bk->j', *x, *y, backend=backend)
+        z2 = contract('bi,bj,bk,bi,bj,bk->k', *x, *y, backend=backend)
+        z = torch.stack([z0, z1, z2])
+
+    Among all valid inputs, some computations are polynomial in the sizes of
+    the input tensors and other computations are exponential in the sizes of
+    the input tensors. This function raises :py:class:`NotImplementedError`
+    whenever the computation is exponential.
 
     :param str equation: an einsum equation, optionally with multiple outputs.
     :param torch.Tensor operands: a collection of tensors
@@ -446,6 +480,10 @@ def ubersum(equation, *operands, **kwargs):
         cache.
     :return: a tuple of tensors of requested shape, one entry per output.
     :rtype: tuple
+    :raises ValueError: if tensor sizes mismatch or an output requests a
+        batched dim without that dim's batch dims.
+    :raises NotImplementedError: if contraction would have cost exponential in
+        the size of any input tensor.
     """
     # Extract kwargs.
     cache = kwargs.pop('cache', None)
@@ -487,3 +525,113 @@ def ubersum(equation, *operands, **kwargs):
                 term = term.permute(*map(dims.index, output))
             results.append(term)
     return tuple(results)
+
+
+def _select(tensor, dims, indices):
+    for dim, index in zip(dims, indices):
+        tensor = tensor.select(dim, index)
+    return tensor
+
+
+class _DimFlattener(object):
+    """
+    Object to map batched dims to batches of flat 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):
+        """
+        Converts a batched dim + batch indices to a flattened dim.
+
+        :param str dim: a batched dimension to flatten
+        :param dict indices: a mapping from batch dimension to int
+        :return: a flattened 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_ubersum(equation, *operands, **kwargs):
+    """
+    Naive reference implementation of :func:`ubersum`.
+
+    This implementation should never raise ``NotImplementedError``.
+    This implementation should agree with :func:`ubersum` whenver
+    :func:`ubersum` does not raise ``NotImplementedError``.
+    """
+    # Parse equation, without loss of generality assuming a single output.
+    inputs, outputs = equation.split('->')
+    outputs = outputs.split(',')
+    if len(outputs) > 1:
+        return tuple(naive_ubersum(inputs + '->' + output, *operands, **kwargs)[0]
+                     for output in outputs)
+    output, = outputs
+    inputs = inputs.split(',')
+
+    # Split dims into batch dims, contraction dims, and dims to keep.
+    batch_dims = set(kwargs.pop('batch_dims', ''))
+    if not batch_dims:
+        result = opt_einsum.contract(equation, *operands, backend='pyro.ops.einsum.torch_log')
+        return (result,)
+    output_dims = set(output)
+
+    # Collect sizes of all dimensions.
+    sizes = {}
+    for input_, operand in zip(inputs, operands):
+        for dim, size in zip(input_, operand.shape):
+            old = sizes.setdefault(dim, size)
+            if old != size:
+                raise ValueError(u"Dimension size mismatch at dim '{}': {} vs {}"
+                                 .format(dim, size, old))
+
+    # Compute batch context for each non-batch dim, by convention the
+    # intersection over all batch contexts of tensors in which the dim appears.
+    dim_to_ordinal = {}
+    for dims in map(set, inputs):
+        ordinal = dims & batch_dims
+        for dim in dims - batch_dims:
+            dim_to_ordinal[dim] = dim_to_ordinal.get(dim, ordinal) & ordinal
+    for dim in output_dims - batch_dims:
+        missing_dims = dim_to_ordinal[dim] - output_dims
+        if missing_dims:
+            raise ValueError(u"It is nonsensical to preserve a batched dim without preserving "
+                             u"all of that dim's batch dims, but found '{}' without '{}' in '{}'"
+                             .format(dim, ','.join(missing_dims), equation))
+
+    # Flatten by replicating along batch dimensions.
+    flatten_dim = _DimFlattener(dim_to_ordinal)
+    flat_inputs = []
+    flat_operands = []
+    for input_, operand in zip(inputs, operands):
+        local_dims = [d for d in input_ if d in batch_dims]
+        offsets = [input_.index(d) - len(input_) for d in local_dims]
+        for index in itertools.product(*(range(sizes[d]) for d in local_dims)):
+            flat_inputs.append(''.join(flatten_dim(d, dict(zip(local_dims, index)))
+                                       for d in input_ if d not in batch_dims))
+            flat_operands.append(_select(operand, offsets, index))
+
+    # Defer to unbatched einsum.
+    result = operands[0].new_empty(torch.Size(sizes[d] for d in output))
+    local_dims = [d for d in output if d in batch_dims]
+    offsets = [output.index(d) - len(output) for d in local_dims]
+    for index in itertools.product(*(range(sizes[d]) for d in local_dims)):
+        flat_output = ''.join(flatten_dim(d, dict(zip(local_dims, index)))
+                              for d in output if d not in batch_dims)
+        flat_equation = ','.join(flat_inputs) + '->' + flat_output
+        flat_result = opt_einsum.contract(flat_equation, *flat_operands,
+                                          backend='pyro.ops.einsum.torch_log')
+        _select(result, offsets, index).copy_(flat_result)
+    return (result,)
diff --git a/pyro/ops/einsum/__init__.py b/pyro/ops/einsum/__init__.py
--- a/pyro/ops/einsum/__init__.py
+++ b/pyro/ops/einsum/__init__.py
@@ -11,7 +11,7 @@
 
 def contract(equation, *operands, **kwargs):
     """
-    Wrpper around :func:`opt_einsum.contract` that caches contraction paths.
+    Wrapper around :func:`opt_einsum.contract` that caches contraction paths.
 
     :param bool cache_path: whether to cache the contraction path.
         Defaults to True.
diff --git a/pyro/ops/einsum/torch_log.py b/pyro/ops/einsum/torch_log.py
--- a/pyro/ops/einsum/torch_log.py
+++ b/pyro/ops/einsum/torch_log.py
@@ -32,11 +32,12 @@ def einsum(equation, *operands):
         exp_operands.append((operand - shift).exp())
 
         # permute shift to match output
-        shift = shift.squeeze()
-        dims = [dim for dim in dims if dim in output]
-        dims = [dim for dim in output if dim not in dims] + dims
-        shift = shift.reshape((1,) * (len(output) - len(shift.shape)) + shift.shape)
-        if dims:
+        shift = shift.reshape(torch.Size(size for size, dim in zip(operand.shape, dims)
+                                         if dim in output))
+        if shift.dim():
+            shift = shift.reshape((1,) * (len(output) - shift.dim()) + shift.shape)
+            dims = [dim for dim in dims if dim in output]
+            dims = [dim for dim in output if dim not in dims] + dims
             shift = shift.permute(*(dims.index(dim) for dim in output))
         shifts.append(shift)
 


Test patch:
diff --git a/tests/infer/test_enum.py b/tests/infer/test_enum.py
--- 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/einsum/test_torch_log.py b/tests/ops/einsum/test_torch_log.py
--- a/tests/ops/einsum/test_torch_log.py
+++ b/tests/ops/einsum/test_torch_log.py
@@ -1,5 +1,7 @@
 from __future__ import absolute_import, division, print_function
 
+import itertools
+
 import pytest
 import torch
 
@@ -9,6 +11,7 @@
 from tests.common import assert_equal
 
 
+@pytest.mark.parametrize('min_size', [1, 2])
 @pytest.mark.parametrize('equation', [
     ',ab->ab',
     'ab,,bc->a',
@@ -30,12 +33,13 @@
     'ab,bc,cd->ac',
     'ab,bc,cd->ad',
     'ab,bc,cd->bc',
+    'a,a,ab,b,b,b,b->a',
 ])
-def test_einsum(equation):
+def test_einsum(equation, min_size):
     inputs, output = equation.split('->')
     inputs = inputs.split(',')
     symbols = sorted(set(equation) - set(',->'))
-    sizes = dict(zip(symbols, range(2, 2 + len(symbols))))
+    sizes = dict(zip(symbols, itertools.count(min_size)))
     shapes = [torch.Size(tuple(sizes[dim] for dim in dims))
               for dims in inputs]
     operands = [torch.randn(shape) for shape in shapes]
diff --git a/tests/ops/test_contract.py b/tests/ops/test_contract.py
--- a/tests/ops/test_contract.py
+++ b/tests/ops/test_contract.py
@@ -1,5 +1,6 @@
 from __future__ import absolute_import, division, print_function
 
+import itertools
 import numbers
 from collections import OrderedDict
 
@@ -8,9 +9,11 @@
 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
+from pyro.util import optional
+from tests.common import assert_equal, xfail_param
 
 
 def deep_copy(x):
@@ -249,50 +252,86 @@ def test_contract_to_tensor_sizes(a, b, c, d):
 
 UBERSUM_EXAMPLES = [
     ('->', ''),
-    ('a->', ''),
-    ('ab->', ''),
-    ('ab,bc->ac', ''),
-    ('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'),
+    ('a->,a', ''),
+    ('ab->,a,b,ab,ba', ''),
+    ('ab,bc->,a,b,c,ab,bc,ac,abc', ''),
+    ('ab,bc,cd->,a,b,c,d,ab,ac,ad,bc,bd,cd,abc,acd,bcd,abcd', ''),
+    ('a->,a', 'a'),
+    (',a->,a', 'a'),
+    (',a,a->,a', 'a'),
+    (',a,ab->,a,ab', 'a'),
+    (',a,a,ab,ab->,a,ab', 'a'),
+    ('ca,ab->,a,ab,ac,abc', 'a'),
+    ('ac,bc,abc->,c,a,ac,b,bc,abc', 'ab'),
+    ('a,bd,abcd->,a,b,bd,ab,abc,abd,abcd', 'ab'),
+    ('ac,bd,abcd->,a,ac,b,bd,ab,abc,abd,abcd', 'ab'),
+    (',a,b,c,ab,ac,bc,abc->,c,a,b,ac,bc,ab,abc', 'ab'),
+    (',ad,abd,acd->,a,b,c,ab,ac,abc,ad,abd,acd,abcd', 'abc'),
 ]
 
 
-@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)))}
+    sizes = {dim: size for dim, size in zip(symbols, itertools.cycle([2, 3, 4]))}
     inputs, outputs = equation.split('->')
+    inputs = inputs.split(',')
+    outputs = outputs.split(',')
     operands = []
-    for dims in inputs.split(','):
+    for dims in inputs:
         shape = tuple(sizes[dim] for dim in dims)
         operands.append(torch.randn(shape))
+    return inputs, outputs, operands, sizes
 
-    actual = ubersum(equation, *operands, batch_dims=batch_dims)
 
-    outputs = outputs.split(',')
+@pytest.mark.parametrize('equation,batch_dims', UBERSUM_EXAMPLES)
+def test_naive_ubersum(equation, batch_dims):
+    inputs, outputs, operands, sizes = make_example(equation)
+
+    actual = naive_ubersum(equation, *operands, batch_dims=batch_dims)
+
+    assert isinstance(actual, tuple)
     assert len(actual) == len(outputs)
     for output, actual_part in zip(outputs, actual):
         expected_shape = tuple(sizes[dim] for dim in output)
         assert actual_part.shape == expected_shape
-        if set(batch_dims) <= set(output):
-            equation_part = inputs + '->' + output
+        if not batch_dims:
+            equation_part = ','.join(inputs) + '->' + output
             expected_part = opt_einsum.contract(equation_part, *operands,
                                                 backend='pyro.ops.einsum.torch_log')
-            assert_equal(actual_part, expected_part)
+            assert_equal(expected_part, actual_part,
+                         msg=u"For output '{}':\nExpected:\n{}\nActual:\n{}".format(
+                             output, expected_part.detach().cpu(), actual_part.detach().cpu()))
+
+
+@pytest.mark.xfail(reason='improper handling of batched output')
+@pytest.mark.parametrize('equation,batch_dims', UBERSUM_EXAMPLES)
+def test_ubersum(equation, batch_dims):
+    inputs, outputs, operands, sizes = make_example(equation)
+
+    try:
+        actual = ubersum(equation, *operands, batch_dims=batch_dims)
+    except NotImplementedError:
+        pytest.skip()
+
+    assert isinstance(actual, tuple)
+    assert len(actual) == len(outputs)
+    expected = naive_ubersum(equation, *operands, batch_dims=batch_dims)
+    for output, expected_part, actual_part in zip(outputs, expected, actual):
+        assert_equal(expected_part, actual_part,
+                     msg=u"For output '{}':\nExpected:\n{}\nActual:\n{}".format(
+                         output, expected_part.detach().cpu(), actual_part.detach().cpu()))
 
 
 @pytest.mark.parametrize('a', [2, 1])
 @pytest.mark.parametrize('b', [3, 1])
 @pytest.mark.parametrize('c', [3, 1])
 @pytest.mark.parametrize('d', [4, 1])
-def test_ubersum_sizes(a, b, c, d):
+@pytest.mark.parametrize('impl', [naive_ubersum, ubersum])
+def test_ubersum_sizes(impl, a, b, c, d):
     X = torch.randn(a, b)
     Y = torch.randn(b, c)
     Z = torch.randn(c, d)
-    actual = ubersum('ab,bc,cd->a,b,c,d', X, Y, Z, batch_dims='ad')
+    actual = impl('ab,bc,cd->a,b,c,d', X, Y, Z, batch_dims='ad')
     actual_a, actual_b, actual_c, actual_d = actual
     assert actual_a.shape == (a,)
     assert actual_b.shape == (b,)
@@ -300,7 +339,8 @@ def test_ubersum_sizes(a, b, c, d):
     assert actual_d.shape == (d,)
 
 
-def test_ubersum_1():
+@pytest.mark.parametrize('impl', [naive_ubersum, ubersum])
+def test_ubersum_1(impl):
     # y {a}   z {b}
     #      \  /
     #     x {}  <--- target
@@ -308,12 +348,13 @@ def test_ubersum_1():
     x = torch.randn(c)
     y = torch.randn(c, d, a)
     z = torch.randn(e, c, b)
-    actual, = ubersum('c,cda,ecb->', x, y, z, batch_dims='ab')
+    actual, = impl('c,cda,ecb->', x, y, z, batch_dims='ab')
     expected = logsumexp(x + logsumexp(y, -2).sum(-1) + logsumexp(z, -3).sum(-1), -1)
     assert_equal(actual, expected)
 
 
-def test_ubersum_2():
+@pytest.mark.parametrize('impl', [naive_ubersum, ubersum])
+def test_ubersum_2(impl):
     # y {a}   z {b}  <--- target
     #      \  /
     #     x {}
@@ -321,13 +362,14 @@ def test_ubersum_2():
     x = torch.randn(c)
     y = torch.randn(c, d, a)
     z = torch.randn(e, c, b)
-    actual, = ubersum('c,cda,ecb->b', x, y, z, batch_dims='ab')
+    actual, = impl('c,cda,ecb->b', x, y, z, batch_dims='ab')
     xyz = logsumexp(x + logsumexp(y, -2).sum(-1) + logsumexp(z, -3).sum(-1), -1)
     expected = xyz.expand(b)
     assert_equal(actual, expected)
 
 
-def test_ubersum_3():
+@pytest.mark.parametrize('impl', [naive_ubersum, ubersum])
+def test_ubersum_3(impl):
     #       z {b,c}
     #           |
     # w {a}  y {b}  <--- target
@@ -338,7 +380,7 @@ def test_ubersum_3():
     x = torch.randn(d)
     y = torch.randn(b, d)
     z = torch.randn(b, c, d, e)
-    actual, = ubersum('ae,d,bd,bcde->be', w, x, y, z, batch_dims='abc')
+    actual, = impl('ae,d,bd,bcde->be', w, x, y, z, batch_dims='abc')
     yz = y.reshape(b, d, 1) + z.sum(-3)  # eliminate c
     assert yz.shape == (b, d, e)
     yz = yz.sum(0)  # eliminate b
@@ -351,7 +393,32 @@ def test_ubersum_3():
     assert_equal(actual, expected)
 
 
-def test_ubersum_collide_error():
+@pytest.mark.parametrize('impl', [
+    naive_ubersum,
+    xfail_param(ubersum, reason='incorrect forward-backward implementation'),
+])
+def test_ubersum_4(impl):
+    # x,y {b}  <--- target
+    #      |
+    #     {}
+    a, b, c, d = 2, 3, 4, 5
+    x = torch.randn(a, b)
+    y = torch.randn(d, b, c)
+    actual, = impl('ab,dbc->dc', x, y, batch_dims='d')
+    x_b1 = logsumexp(x, 0).unsqueeze(-1)
+    assert x_b1.shape == (b, 1)
+    y_db1 = logsumexp(y, 2, keepdim=True)
+    assert y_db1.shape == (d, b, 1)
+    y_dbc = y_db1.sum(0) - y_db1 + y  # avoid double counting
+    assert y_dbc.shape == (d, b, c)
+    xy_dc = logsumexp(x_b1 + y_dbc, 1)
+    assert xy_dc.shape == (d, c)
+    expected = xy_dc
+    assert_equal(actual, expected)
+
+
+@pytest.mark.parametrize('impl,implemented', [(naive_ubersum, True), (ubersum, False)])
+def test_ubersum_collide_implemented(impl, implemented):
     # Non-tree iaranges cause exponential blowup,
     # so ubersum() refuses to evaluate them.
     #
@@ -364,11 +431,13 @@ 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'):
-        ubersum('ac,bd,abcd->', x, y, z, batch_dims='ab')
+    raises = pytest.raises(NotImplementedError, match='Expected tree-structured iarange nesting')
+    with optional(raises, not implemented):
+        impl('ac,bd,abcd->', x, y, z, batch_dims='ab')
 
 
-def test_ubersum_collide_ok_1():
+@pytest.mark.parametrize('impl', [naive_ubersum, ubersum])
+def test_ubersum_collide_ok_1(impl):
     # The following is ok because it splits into connected components
     # {x,z1} and {y,z2}, thereby avoiding exponential blowup.
     #
@@ -382,10 +451,11 @@ def test_ubersum_collide_ok_1():
     y = torch.randn(b, d)
     z1 = torch.randn(a, b, c)
     z2 = torch.randn(a, b, d)
-    ubersum('ac,bd,abc,abd->', x, y, z1, z2, batch_dims='ab')
+    impl('ac,bd,abc,abd->', x, y, z1, z2, batch_dims='ab')
 
 
-def test_ubersum_collide_ok_2():
+@pytest.mark.parametrize('impl', [naive_ubersum, ubersum])
+def test_ubersum_collide_ok_2(impl):
     # The following is ok because z1 can be contracted to x and
     # z2 can be contracted to y.
     #
@@ -400,10 +470,11 @@ def test_ubersum_collide_ok_2():
     y = torch.randn(b, d)
     z1 = torch.randn(a, b, c)
     z2 = torch.randn(a, b, d)
-    ubersum('cd,ac,bd,abc,abd->', w, x, y, z1, z2, batch_dims='ab')
+    impl('cd,ac,bd,abc,abd->', w, x, y, z1, z2, batch_dims='ab')
 
 
-def test_ubersum_collide_ok_3():
+@pytest.mark.parametrize('impl', [naive_ubersum, ubersum])
+def test_ubersum_collide_ok_3(impl):
     # The following is ok because x, y, and z can be independently contracted to w.
     #
     #      z {a,b}
@@ -416,17 +487,41 @@ def test_ubersum_collide_ok_3():
     x = torch.randn(a, c)
     y = torch.randn(b, c)
     z = torch.randn(a, b, c)
-    ubersum('c,ac,bc,abc->', w, x, y, z, batch_dims='ab')
+    impl('c,ac,bc,abc->', w, x, y, z, batch_dims='ab')
 
 
-UBERSUM_ERRORS = [
+UBERSUM_SHAPE_ERRORS = [
     ('ab,bc->', [(2, 3), (4, 5)], ''),
     ('ab,bc->', [(2, 3), (4, 5)], 'b'),
 ]
 
 
-@pytest.mark.parametrize('equation,shapes,batch_dims', UBERSUM_ERRORS)
-def test_ubersum_size_error(equation, shapes, batch_dims):
+@pytest.mark.parametrize('equation,shapes,batch_dims', UBERSUM_SHAPE_ERRORS)
+@pytest.mark.parametrize('impl', [naive_ubersum, ubersum])
+def test_ubersum_size_error(impl, equation, shapes, batch_dims):
     operands = [torch.randn(shape) for shape in shapes]
-    with pytest.raises(ValueError, match='Dimension size mismatch'):
-        ubersum(equation, *operands, batch_dims=batch_dims)
+    with pytest.raises(ValueError, match='Dimension size mismatch|Size of label'):
+        impl(equation, *operands, batch_dims=batch_dims)
+
+
+UBERSUM_BATCH_ERRORS = [
+    ('ab->b', 'a'),
+    (',ab->b', 'a'),
+    ('ac,abc->c', 'a'),
+    (',ac,abc->c', 'a'),
+    ('abc->ac', 'ab'),
+    ('abc->bc', 'ab'),
+]
+
+
+@pytest.mark.parametrize('equation,batch_dims', UBERSUM_BATCH_ERRORS)
+@pytest.mark.parametrize('impl', [
+    naive_ubersum,
+    xfail_param(ubersum, reason='does not detect senseless batch output'),
+])
+def test_ubersum_batch_error(impl, equation, batch_dims):
+    inputs, outputs = equation.split('->')
+    operands = [torch.randn(torch.Size((2,) * len(input_)))
+                for input_ in inputs.split(',')]
+    with pytest.raises(ValueError, match='It is nonsensical to preserve a batched dim'):
+        impl(equation, *operands, batch_dims=batch_dims)