diff --git a/finat/physically_mapped.py b/finat/physically_mapped.py index 511d3e14..da88e7d6 100644 --- a/finat/physically_mapped.py +++ b/finat/physically_mapped.py @@ -1,5 +1,6 @@ from abc import ABCMeta, abstractmethod -from collections.abc import Mapping +from collections.abc import Iterable, Mapping +from functools import cached_property import gem import numpy @@ -7,6 +8,10 @@ from finat.citations import cite +zero = gem.Zero() +one = gem.Literal(1.0) + + class NeedsCoordinateMappingElement(metaclass=ABCMeta): """Abstract class for elements that require physical information either to map or construct their basis functions.""" @@ -16,44 +21,118 @@ def dual_transformation(self, Q, coordinate_mapping=None): class MappedTabulation(Mapping): - """A lazy tabulation dict that applies the basis transformation only - on the requested derivatives. + """Apply a sparse basis transformation to reference tabulations. + + Parameters + ---------- + M : gem.ListTensor + Basis-transformation matrix. + ref_tabulation : Mapping + Reference tabulations indexed by derivative order. + indices : iterable of int, optional + Rows retained by an element restriction. + + Notes + ----- + In order to generate good loopy kernels, rows are padded so that they have + the same number of entries. Constant tables select the reference column + and one of the distinct symbolic coefficients. Interning coefficients + preserves their sharing without materialising a symbolic matrix entry by + entry. - :arg M: a gem.ListTensor with the basis transformation matrix. - :arg ref_tabulation: a dict of tabulations on the reference cell. - :kwarg indices: an optional list of restriction indices on the basis functions. """ - def __init__(self, M, ref_tabulation, indices=None): - self.M = M + + def __init__( + self, M: gem.ListTensor, ref_tabulation: Mapping, + indices: Iterable[int] | None = None) -> None: self.ref_tabulation = ref_tabulation if indices is None: - indices = list(range(M.shape[0])) - self.indices = indices - # we expect M to be sparse with O(1) nonzeros per row - # for each row, get the column index of each nonzero entry - csr = [[j for j in range(M.shape[1]) if not isinstance(M.array[i, j], gem.Zero)] - for i in indices] - self.csr = csr + indices = range(M.shape[0]) + self.indices = tuple(indices) + self._space_dim = len(self.indices) + self._value_dim = M.shape[1] + + nonzero_rows = [] + for source_row in self.indices: + row = [] + for column in range(M.shape[1]): + value = M.array[source_row, column] + if not isinstance(value, gem.Zero): + row.append((column, value)) + nonzero_rows.append(row) + width = max((len(row) for row in nonzero_rows), default=0) + nrows = len(self.indices) + columns = numpy.zeros((nrows, width), dtype=gem.uint_type) + data = numpy.full((nrows, width), zero, dtype=object) + for index, row in enumerate(nonzero_rows): + columns[index, :len(row)] = tuple(column for column, _ in row) + data[index, :len(row)] = tuple(gem.as_gem(value) for _, value in row) + self._width = width + self._columns = gem.Literal(columns, dtype=gem.uint_type) + values = [] + value_numbers = {} + value_indices = numpy.empty(data.shape, dtype=gem.uint_type) + for multiindex, value in numpy.ndenumerate(data): + try: + number = value_numbers[value] + except KeyError: + number = len(values) + value_numbers[value] = number + values.append(value) + value_indices[multiindex] = number + self._value_indices = gem.Literal(value_indices, dtype=gem.uint_type) + self._values = gem.ListTensor(values) self._tabulation_cache = {} - def matvec(self, table): - # basis recombination using hand-rolled sparse-dense matrix multiplication - ii = gem.indices(len(table.shape)-1) - phi = [gem.Indexed(table, (j, *ii)) for j in range(self.M.shape[1])] - # the sum approach is faster than calling numpy.dot or gem.IndexSum - exprs = [gem.ComponentTensor(gem.Sum(*(self.M.array[i, j] * phi[j] for j in js)), ii) - for i, js in zip(self.indices, self.csr)] + @cached_property + def _reference_index(self) -> gem.Index: + """Contraction over the reference basis, shared by all tabulations.""" + return gem.Index(extent=self._value_dim) + + @cached_property + def _row_index(self) -> gem.Index: + """Contraction over a padded row, shared by all tabulations.""" + return gem.Index(extent=self._width) + + def _entry(self, r: gem.Index, a: gem.Index) -> gem.Node: + """Entry ``M[r, a]`` of the basis transformation. + + Parameters + ---------- + r + Index over the rows retained by the element. + a + Index over the reference basis. + + Returns + ------- + gem.Node + A sum over the padded row of an interned entry against a Delta + selecting its column, so that contracting either axis of ``M`` + is ordinary GEM algebra. - result = gem.ListTensor(exprs) - result, = gem.optimise.unroll_indexsum((result,), lambda index: True) - # result = gem.optimise.aggressive_unroll(self.M @ table) - return result + """ + k = self._row_index + entry = gem.Indexed( + self._values, + (gem.VariableIndex(gem.Indexed(self._value_indices, (r, k))),)) + column = gem.VariableIndex(gem.Indexed(self._columns, (r, k))) + return gem.IndexSum(gem.Product(entry, gem.Delta(column, a)), (k,)) + + def matmul(self, table: gem.Node) -> gem.Node: + """Apply the basis transformation to a reference tabulation.""" + r = gem.Index(extent=self._space_dim) + a = self._reference_index + tail = gem.indices(len(table.shape) - 1) + mapped = gem.IndexSum( + gem.Product(self._entry(r, a), gem.Indexed(table, (a, *tail))), (a,)) + return gem.ComponentTensor(mapped, (r, *tail)) def __getitem__(self, alpha): try: return self._tabulation_cache[alpha] except KeyError: - result = self.matvec(self.ref_tabulation[alpha]) + result = self.matmul(self.ref_tabulation[alpha]) return self._tabulation_cache.setdefault(alpha, result) def __iter__(self): @@ -195,10 +274,6 @@ def physical_vertices(self): (gdim, ).""" -zero = gem.Zero() -one = gem.Literal(1.0) - - def identity(*shape): V = numpy.eye(*shape, dtype=object) for multiindex in numpy.ndindex(V.shape): diff --git a/gem/cost.py b/gem/cost.py index a46a2238..7451b576 100644 --- a/gem/cost.py +++ b/gem/cost.py @@ -86,6 +86,18 @@ def iteration_count(indices: Iterable[Index]) -> int: return int(numpy.prod([index.extent for index in indices], dtype=int)) +def index_space_literal(indices: Iterable[Index]) -> Literal: + """The cardinality of a rectangular index space, as a scalar. + + The empty product is one, so contracting it counts the tuples in the + index space. + + :arg indices: indices spanning the space + :returns: the number of points, as a floating point literal + """ + return Literal(float(iteration_count(indices))) + + def operation_count(node: Node) -> int: """Estimate the scalar operations performed by one GEM node. diff --git a/gem/gem.py b/gem/gem.py index 46c6d9b1..2dc0fc80 100644 --- a/gem/gem.py +++ b/gem/gem.py @@ -681,6 +681,15 @@ def __reduce__(self): return type(self), (self.expression,) +def _index_free_indices(index): + """Return the free indices represented by an index expression.""" + if isinstance(index, Index): + return (index,) + if isinstance(index, VariableIndex): + return index.expression.free_indices + return () + + class Indexed(Scalar): __slots__ = ('children', 'multiindex', 'indirect_children') __back__ = ('multiindex',) @@ -1072,9 +1081,11 @@ def __new__(cls, i, j, dtype=None): self = super(Delta, cls).__new__(cls) self.i = i self.j = j - # Set up free indices - free_indices = [index for index in (i, j) if isinstance(index, Index)] - self.free_indices = tuple(unique(free_indices)) + # Set up free indices. A VariableIndex is not itself a free index, + # but the expression it wraps may be free in others; those propagate + # here exactly as they do through Indexed. + self.free_indices = tuple(unique(chain.from_iterable( + _index_free_indices(index) for index in (i, j)))) self._dtype = dtype return self diff --git a/gem/interpreter.py b/gem/interpreter.py index 13eeb44a..b2dd609e 100644 --- a/gem/interpreter.py +++ b/gem/interpreter.py @@ -263,8 +263,34 @@ def _evaluate_conditional(e, self): def _evaluate_indexed(e, self): """Indexing maps shape to free indices""" val = self(e.children[0]) - fids = tuple(i for i in e.multiindex if isinstance(i, gem.Index)) + variable_indices = {i: self(i.expression) for i in e.multiindex + if isinstance(i, gem.VariableIndex)} + + if any(result.fids for result in variable_indices.values()): + # Some variable index depends on free indices: gather entries + # one by one over the extent of the free indices. + fids = list(val.fids) + for i in e.multiindex: + new_fids = (i,) if isinstance(i, gem.Index) else \ + variable_indices[i].fids if isinstance(i, gem.VariableIndex) else () + fids.extend(f for f in new_fids if f not in fids) + fids = tuple(fids) + out = numpy.empty(tuple(f.extent for f in fids), dtype=val.arr.dtype) + for idx in numpy.ndindex(out.shape): + env = dict(zip(fids, idx)) + vidx = [env[f] for f in val.fids] + for i in e.multiindex: + if isinstance(i, gem.Index): + vidx.append(env[i]) + elif isinstance(i, gem.VariableIndex): + result = variable_indices[i] + vidx.append(int(result.arr[tuple(env[f] for f in result.fids)])) + else: + vidx.append(i) + out[idx] = val.arr[tuple(vidx)] + return Result(out, fids) + fids = tuple(i for i in e.multiindex if isinstance(i, gem.Index)) idx = [] # First pick up all the existing free indices for _ in val.fids: @@ -275,10 +301,10 @@ def _evaluate_indexed(e, self): # Free index, want entire extent idx.append(slice(None)) elif isinstance(i, gem.VariableIndex): - # Variable index, evaluate inner expression - result, = self(i.expression) + # Variable index, constant during kernel execution + result = variable_indices[i] assert not result.tshape - idx.append(result[()]) + idx.append(int(result.arr[()])) else: # Fixed index, just pick that value idx.append(i) diff --git a/gem/node.py b/gem/node.py index 190fe6d4..34e05775 100644 --- a/gem/node.py +++ b/gem/node.py @@ -103,7 +103,12 @@ def get_hash(self): return hash((type(self), *self._arguments)) -def _make_traversal_children(node): +def traversal_children(node): + """The children a DAG walk descends into, index expressions included. + + :arg node: a GEM expression + :returns: the child nodes, plus the nodes hidden in index expressions + """ if isinstance(node, (gem.Indexed, gem.FlexiblyIndexed)): # Include child nodes hidden in index expressions. return node.children + node.indirect_children @@ -117,7 +122,7 @@ def pre_traversal(expression_dags): Notes ----- This function also walks through nodes in index expressions - (e.g., `VariableIndex`s); see ``_make_traversal_children()``. + (e.g., `VariableIndex`s); see ``traversal_children()``. """ seen = set() @@ -133,7 +138,7 @@ def pre_traversal(expression_dags): while lifo: node = lifo.pop() yield node - children = _make_traversal_children(node) + children = traversal_children(node) for child in reversed(children): if child not in seen: seen.add(child) @@ -146,7 +151,7 @@ def post_traversal(expression_dags): Notes ----- This function also walks through nodes in index expressions - (e.g., `VariableIndex`s); see ``_make_traversal_children()``. + (e.g., `VariableIndex`s); see ``traversal_children()``. """ @@ -158,13 +163,13 @@ def post_traversal(expression_dags): for root in expression_dags: if root not in seen: seen.add(root) - lifo.append((root, list(_make_traversal_children(root)))) + lifo.append((root, list(traversal_children(root)))) while lifo: node, deps = lifo[-1] for i, dep in enumerate(deps): if dep is not None and dep not in seen: - lifo.append((dep, list(_make_traversal_children(dep)))) + lifo.append((dep, list(traversal_children(dep)))) deps[i] = None break else: @@ -184,12 +189,12 @@ def collect_refcount(expression_dags): ----- This function also collects reference counts of nodes in index expressions (e.g., `VariableIndex`s); see - ``_make_traversal_children()``. + ``traversal_children()``. """ result = collections.Counter(expression_dags) for node in traversal(expression_dags): - result.update(_make_traversal_children(node)) + result.update(traversal_children(node)) return result diff --git a/gem/optimise.py b/gem/optimise.py index e367d4d2..e5afdd98 100644 --- a/gem/optimise.py +++ b/gem/optimise.py @@ -9,12 +9,13 @@ import numpy +from gem.cost import index_space_literal from gem.utils import groupby from gem.node import (Memoizer, MemoizerArg, reuse_if_untouched, - reuse_if_untouched_arg, traversal) + reuse_if_untouched_arg, traversal, traversal_children) from gem.gem import (Node, Failure, Identity, Constant, Literal, Zero, Product, Sum, Comparison, Conditional, Division, - Index, VariableIndex, Indexed, FlexiblyIndexed, + Index, IndexBase, VariableIndex, Indexed, FlexiblyIndexed, IndexSum, ComponentTensor, ListTensor, Delta, partial_indexed, one) @@ -100,7 +101,13 @@ def replace_indices(node, self, subst): def _replace_indices_atomic(i, self, subst): if isinstance(i, VariableIndex): new_expr = self(i.expression, subst) - return i if new_expr == i.expression else VariableIndex(new_expr) + if new_expr == i.expression: + return i + # A variable index that substitution has made constant is a fixed + # index, and folding it lets the lookup itself be evaluated. + if isinstance(new_expr, Literal) and not new_expr.shape: + return int(new_expr.array) + return VariableIndex(new_expr) else: substitute = dict(subst) return substitute.get(i, i) @@ -270,6 +277,15 @@ def child(expression): elif all(e.j == k and e.i == expr.i for k, e in enumerate(expressions)): return expr.reconstruct(expr.i, index) + if types == {IndexSum}: + extents = {tuple(i.extent for i in e.multiindex) for e in expressions} + if len(extents) == 1: + multiindex = tuple(Index(extent=extent) for extent in extents.pop()) + summands = [Indexed(ComponentTensor(e.children[0], e.multiindex), + multiindex) + for e in expressions] + return IndexSum(_select_expression(summands, index), multiindex) + if len(types) == 1: cls, = types if cls.__front__ or cls.__back__: @@ -309,12 +325,15 @@ def select_expression(expressions, index): return ComponentTensor(selected, alpha) -def delta_elimination(sum_indices, factors, index_replacer=None): +def delta_elimination(sum_indices, factors, index_replacer=None, indirect_only=False): """IndexSum-Delta cancellation. :arg sum_indices: free indices for contractions :arg factors: product factors :kwarg index_replacer: MemoizerArg(filtered_replace_indices) + :kwarg indirect_only: only cancel a Delta that compares a + :class:`~.VariableIndex`, leaving one between two + plain indices to a later pass :returns: optimised (sum_indices, factors) """ @@ -331,9 +350,15 @@ def substitute(expression, from_, to_): else: return Indexed(ComponentTensor(expression, (from_,)), (to_,)) - delta_queue = [(f, index) - for f in factors if isinstance(f, Delta) - for index in (f.i, f.j) if index in sum_indices] + def cancellable(factors): + return [(f, index) + for f in factors if isinstance(f, Delta) + if not indirect_only + or isinstance(f.i, VariableIndex) or isinstance(f.j, VariableIndex) + for index in (f.i, f.j) + if index in sum_indices] + + delta_queue = cancellable(factors) while delta_queue: delta, from_ = delta_queue[0] to_, = list({delta.i, delta.j} - {from_}) @@ -342,9 +367,7 @@ def substitute(expression, from_, to_): factors = [substitute(f, from_, to_) for f in factors] - delta_queue = [(f, index) - for f in factors if isinstance(f, Delta) - for index in (f.i, f.j) if index in sum_indices] + delta_queue = cancellable(factors) return sum_indices, factors @@ -583,6 +606,9 @@ def _sum_factorise_connected(sum_indices, groups): :arg groups: product factors, grouped by free indices :returns: optimised GEM expression """ + if not groups: + return index_space_literal(sum_indices) + if len(groups) <= _MAX_PLANNED_FACTORS: return _plan_contraction(sum_indices, groups) @@ -611,9 +637,8 @@ def sum_factorise(sum_indices, factors): :arg factors: product factors :returns: optimised GEM expression """ - if len(factors) == 0 and len(sum_indices) == 0: - # Empty product - return one + if len(factors) == 0: + return index_space_literal(sum_indices) # Form groups by free indices groups = groupby(factors, key=lambda f: f.free_indices) @@ -954,6 +979,157 @@ def repeated_contractions(expression): return frozenset(expr for expr, count in counts.items() if count > 1) +def _indirect_delta_axes(node: Node, self: Memoizer) -> frozenset: + """The axes compared by the indirect Deltas below a node, including its own. + + An indirect Delta compares a :class:`~.VariableIndex`, and is the only + kind `cancel_nested_deltas` cancels. Memoising this over the DAG keeps + the search for one linear, rather than re-walking the subtree at every + enclosing contraction. + + :arg node: a GEM expression + :arg self: memoizer visiting the DAG + :returns: the indices some indirect Delta at or below ``node`` compares + """ + axes = frozenset().union(*map(self, traversal_children(node))) + if isinstance(node, Delta) and any(isinstance(i, VariableIndex) + for i in (node.i, node.j)): + axes = axes | {node.i, node.j} + return axes + + +def _constant_map(index: IndexBase) -> tuple | None: + """The literal table behind a VariableIndex, and the indices addressing it. + + Parameters + ---------- + index + Index to inspect. + + Returns + ------- + tuple or None + ``(array, indices)`` when the index is a lookup into a Literal with a + plain multiindex, otherwise None. + + """ + if not isinstance(index, VariableIndex): + return None + expression = index.expression + if not isinstance(expression, Indexed): + return None + table, = expression.children + if not isinstance(table, Literal): + return None + if not all(isinstance(i, Index) for i in expression.multiindex): + return None + return table.array, expression.multiindex + + +def pull_back_indirect_delta( + sum_indices: Iterable[Index], + factors: Iterable[Node], + replacer: MemoizerArg) -> tuple: + """Contract an indirect Delta's own axes before its column axis. + + ``sum_a (sum_rk v(r,k) delta(c(r,k), a)) T(a, q)`` is cancelled by + `delta_elimination` substituting ``a := c(r,k)``, which makes ``T`` + depend on ``r`` and ``k`` and so forces that contraction inside the ``q`` + loop. When ``r`` and ``k`` are contracted here and ``T`` carries indices + of its own, summing them first is cheaper: it yields a dense vector + indexed by ``a``. Run this before `delta_elimination` to take that + cheaper route where it exists. + + :arg sum_indices: indices contracted over the product + :arg factors: product factors + :arg replacer: ``MemoizerArg(filtered_replace_indices)`` + :returns: new ``(sum_indices, factors)``, unchanged when no Delta is + worth pulling back + """ + for delta in factors: + if not isinstance(delta, Delta): + continue + column = delta.j if isinstance(delta.i, VariableIndex) else delta.i + variable = delta.i if isinstance(delta.i, VariableIndex) else delta.j + if not isinstance(column, Index) or not isinstance(variable, VariableIndex): + continue + lookup = _constant_map(variable) + if lookup is None: + continue + table, source_indices = lookup + sources = frozenset(source_indices) + if column not in sum_indices or not sources <= set(sum_indices): + continue + + others = [f for f in factors if f is not delta] + spanning = [f for f in others if column in f.free_indices] + pulled = [f for f in others if column not in f.free_indices] + if not spanning or not pulled: + continue + # Cancelling couples the spanning factors to the source indices. That + # only costs anything when they carry indices of their own. + if not any(set(f.free_indices) - sources - {column} for f in spanning): + continue + + vector = numpy.empty(column.extent, dtype=object) + contributions = defaultdict(list) + for position in numpy.ndindex(table.shape): + substitution = tuple(zip(source_indices, (int(p) for p in position))) + contributions[int(table[position])].append(substitution) + for value in range(column.extent): + terms = [make_product([replacer(f, substitution) for f in pulled]) + for substitution in contributions.get(value, ())] + # A reference basis function that no row maps onto contributes + # nothing. + vector[value] = make_sum(terms) if terms else Zero() + + rest = tuple(i for i in sum_indices if i not in sources) + return rest, [Indexed(ListTensor(vector), (column,)), *spanning] + + return sum_indices, factors + + +def cancel_nested_deltas(expression: Node) -> Node: + """Cancel the indirect Deltas at every contraction of a whole DAG. + + `delta_elimination` only inspects top-level product factors, so a Delta + inside a preserved linear map is invisible to it. Flattening the product + tree first exposes it, and hoists the contractions it sits under so that + substituting the Delta's variable index cannot capture them. + + A Delta comparing a :class:`~.VariableIndex` is the only kind handled + here. It is the only kind that has to be: nothing downstream can lower + one. A Delta between two plain indices is left to monomial collection, + which cancels it knowing what the substitution costs there. + + :arg expression: root of a scalar GEM expression + :returns: the expression with those Deltas cancelled + """ + replacer = MemoizerArg(filtered_replace_indices) + delta_axes = Memoizer(_indirect_delta_axes) + + def visit(node, self): + node = reuse_if_untouched(node, self) + if not isinstance(node, IndexSum): + return node + if not delta_axes(node).intersection(node.multiindex): + return node + sum_indices, factors = traverse_product(node, index_replacer=replacer) + cancelled, new_factors = pull_back_indirect_delta( + sum_indices, factors, replacer) + cancelled, new_factors = delta_elimination( + cancelled, new_factors, index_replacer=replacer, indirect_only=True) + if tuple(cancelled) == tuple(sum_indices) and tuple(new_factors) == tuple(factors): + # Nothing cancelled, so rebuilding would only flatten the + # contractions this node nests into a single product, and sum + # factorisation needs them nested. + return node + factors = [replacer(factor, ()) for factor in new_factors] + return IndexSum(make_product(factors), tuple(cancelled)) + + return Memoizer(visit)(expression) + + def contraction(expression): """Optimise the contractions of the tensor product at the root of the expression, including: @@ -980,7 +1156,10 @@ def rebuild(expression): sum_indices, factors = traverse_product( expression, index_replacer=index_replacer, stop_at=lambda e: e is not root and e in keep) - sum_indices, factors = delta_elimination(sum_indices, factors, index_replacer=index_replacer) + sum_indices, factors = pull_back_indirect_delta( + sum_indices, factors, index_replacer) + sum_indices, factors = delta_elimination( + sum_indices, factors, index_replacer=index_replacer) factors = [index_replacer(f, ()) for f in factors] return sum_factorise(sum_indices, factors) diff --git a/test/finat/test_zany_mapping.py b/test/finat/test_zany_mapping.py index 9220dca9..f0742caf 100644 --- a/test/finat/test_zany_mapping.py +++ b/test/finat/test_zany_mapping.py @@ -1,11 +1,49 @@ import FIAT import finat +import gem import numpy as np import pytest import pprint from gem.interpreter import evaluate -from finat.physically_mapped import PhysicallyMappedElement +from gem.node import traversal +from gem.optimise import contraction +from finat.physically_mapped import MappedTabulation, PhysicallyMappedElement + + +def test_sparse_mapped_tabulation(): + """Apply a sparse basis map at the cost of its nonzeros.""" + coefficient = gem.Variable("coefficient", ()) + matrix = gem.ListTensor(np.asarray([ + [gem.Literal(1.0), gem.Zero(), coefficient], + [gem.Zero(), gem.Literal(1.0), gem.Zero()], + ], dtype=object)) + table_values = np.arange(1.0, 7.0).reshape(3, 2) + table = gem.Literal(table_values) + + mapped_tabulation = MappedTabulation(matrix, {None: table}) + # Equal symbolic entries share one coefficient slot, so the sparse map + # indexes a vector rather than materialising a symbolic matrix. + assert mapped_tabulation._values.shape == (3,) + + i, j = gem.indices(2) + mapped = contraction(gem.Indexed(mapped_tabulation[None], (i, j))) + + # The three unit entries cost no multiplication, and the one remaining + # nonzero costs exactly one. Nothing is selected by a branch. + products = [node for node in traversal((mapped,)) + if isinstance(node, gem.Product)] + assert len(products) == 1 + assert not any(isinstance(node, gem.Conditional) + for node in traversal((mapped,))) + + actual, = evaluate( + [gem.ComponentTensor(mapped, (i, j))], + {coefficient: np.asarray(2.0)}, + ) + expected = np.asarray([[1.0, 0.0, 2.0], [0.0, 1.0, 0.0]]) \ + @ table_values + assert np.array_equal(actual.arr, expected) def make_unisolvent_points(element, interior=False): diff --git a/test/gem/test_sum_factorise.py b/test/gem/test_sum_factorise.py index ad7e30f8..58df4945 100644 --- a/test/gem/test_sum_factorise.py +++ b/test/gem/test_sum_factorise.py @@ -264,3 +264,20 @@ def test_estimate_cost_counts_the_contraction(): assert flops > 0 assert storage >= largest > 0 assert nodes > 0 + + +def test_empty_product_contraction_counts_index_tuples() -> None: + i, j = gem.Index(extent=2), gem.Index(extent=3) + expression = sum_factorise((i, j), ()) + result, = evaluate([expression]) + assert result.arr == 6 + + +def test_contraction_counts_index_absent_from_factors() -> None: + i, j = gem.Index(extent=2), gem.Index(extent=3) + factor = gem.Indexed(gem.Literal(numpy.arange(3.0)), (j,)) + + expression = sum_factorise((i, j), (factor,)) + result, = evaluate([expression]) + + assert result.arr == 2 * numpy.arange(3.0).sum()