Fix object backend broadcasting and empty reductions

docs/getting_started/backends.md

@@ -337,9 +337,12 @@ There are a few things to note here:
 
 - The returned array is a `numpy.ndarray` but since it has `dtype=object`
   it can really hold *any* python objects
-- We had to explicitly use `backend='object'`, since `numpy.einsum`
-  would have otherwise been dispatched to, which can't handle `dtype=object`
-  (though `numpy.tensordot` in fact can)
+- NumPy's `einsum` supports `dtype=object` starting with NumPy 1.25, but its
+  reductions can require objects to support addition with the Python integer
+  zero. Use `backend='object'` for objects that only support addition and
+  multiplication with each other. Nonempty reductions start with the first
+  product; empty reductions return the Python integer zero.
+- Broadcasting uses NumPy views without copying the input elements.
 - Although an optimized pairwise contraction order is used, the looping in each
   single contraction is **performed in python so performance will be
   drastically lower than for numeric dtypes!**

opt_einsum/backends/object_arrays.py

@@ -33,10 +33,25 @@ def object_einsum(eq: str, *arrays: ArrayType) -> ArrayType:
     lhs, output = eq.split("->")
     inputs = lhs.split(",")
 
-    sizes = {}
-    for term, array in zip(inputs, arrays):
+    sizes: dict[str, int] = {}
+    for operand, (term, array) in enumerate(zip(inputs, arrays)):
+        operand_sizes: dict[str, int] = {}
         for k, d in zip(term, array.shape):
-            sizes[k] = d
+            if k in operand_sizes and operand_sizes[k] != d:
+                raise ValueError(
+                    f"dimensions in operand {operand} for collapsing index '{k}' "
+                    f"don't match ({operand_sizes[k]} != {d})"
+                )
+            operand_sizes[k] = d
+            if k not in sizes or sizes[k] == 1:
+                sizes[k] = d
+
+    # Broadcast views up front so the inner loop can use ordinary coordinates.
+    broadcasted_arrays = []
+    for term, array in zip(inputs, arrays):
+        target_shape = tuple(sizes[k] for k in term)
+        broadcasted_arrays.append(np.broadcast_to(array, target_shape))
+    arrays = tuple(broadcasted_arrays)
 
     out_size = tuple(sizes[k] for k in output)
     out = np.empty(out_size, dtype=object)
@@ -54,6 +69,8 @@ def gen_inner_sum():
                 elements = (array[loc] for array, loc in zip(arrays, locs))
                 yield functools.reduce(operator.mul, elements)
 
-        out[coo_o] = functools.reduce(operator.add, gen_inner_sum())
+        terms = gen_inner_sum()
+        # Empty reductions are zero; nonempty ones need only object addition.
+        out[coo_o] = functools.reduce(operator.add, terms, next(terms, 0))
 
     return out

opt_einsum/tests/test_backends.py

@@ -482,3 +482,68 @@ def test_object_arrays_backend(string: str) -> None:
     obj_opt = expr(*obj_views, backend="object")
     assert obj_opt.dtype == object
     assert np.allclose(ein, obj_opt.astype(float))
+
+
+@pytest.mark.parametrize(
+    "equation, shapes",
+    [
+        ("i,i->i", [(1,), (3,)]),
+        ("i,i->i", [(3,), (1,)]),
+        ("i,i->", [(1,), (3,)]),
+        ("i,i->", [(3,), (1,)]),
+        ("aab,bc->ac", [(2, 2, 1), (3, 2)]),
+        ("...ii,...i->...", [(2, 3, 3), (1, 3)]),
+        ("i,i,i->i", [(0,), (1,), (1,)]),
+        ("i,i,i->", [(1,), (0,), (1,)]),
+        ("i->", [(0,)]),
+        ("ii->", [(0, 0)]),
+        ("ij->i", [(2, 0)]),
+        ("ij,jk->ik", [(2, 0), (1, 3)]),
+    ],
+)
+@pytest.mark.parametrize("use_blas", [False, True])
+def test_object_arrays_broadcasting_and_empty_reductions(equation, shapes, use_blas) -> None:
+    np = pytest.importorskip("numpy")
+    operands = [np.arange(np.prod(shape)).reshape(shape) + 1 for shape in shapes]
+    expected = np.einsum(equation, *operands)
+    objects = [operand.astype(object) for operand in operands]
+
+    actual = contract(equation, *objects, backend="object", use_blas=use_blas)
+    assert actual.dtype == object
+    np.testing.assert_array_equal(actual, expected)
+
+    expr = contract_expression(equation, *shapes, use_blas=use_blas)
+    actual = expr(*objects, backend="object")
+    assert actual.dtype == object
+    np.testing.assert_array_equal(actual, expected)
+
+
+@pytest.mark.parametrize("shape", [(1, 2), (2, 1), (0, 1), (1, 0), (2, 3), (3, 2)])
+@pytest.mark.parametrize("optimize", [False, True])
+def test_object_arrays_reject_mismatched_diagonal(shape, optimize) -> None:
+    np = pytest.importorskip("numpy")
+    operand = np.empty(shape, dtype=object)
+
+    with pytest.raises(ValueError):
+        np.einsum("ii->", operand)
+    with pytest.raises(ValueError):
+        contract("ii->", operand, backend="object", optimize=optimize)
+
+
+@pytest.mark.parametrize("size", [1, 2])
+def test_object_arrays_without_integer_addition(size: int) -> None:
+    np = pytest.importorskip("numpy")
+
+    class Number:
+        def __init__(self, value):
+            self.value = value
+
+        def __mul__(self, other):
+            return Number(self.value * other.value)
+
+        def __add__(self, other):
+            return Number(self.value + other.value)
+
+    values = np.array([Number(i + 1) for i in range(size)], dtype=object)
+    result = contract("i->", values, backend="object")
+    assert result.item().value == sum(range(1, size + 1))