From a92605f0b778bdede7f44b51864e600013c26b7d Mon Sep 17 00:00:00 2001
From: Pablo Brubeck <brubeck@protonmail.com>
Date: Fri, 21 Aug 2026 14:53:49 +0100
Subject: [PATCH 01/11] Normalise scalar factors in shared linear maps

---
 gem/coffee.py                  | 71 +++++++++++++++++++++++++++++++---
 test/gem/test_sum_factorise.py | 28 ++++++++++++++
 2 files changed, 94 insertions(+), 5 deletions(-)

diff --git a/gem/coffee.py b/gem/coffee.py
index 607439fd..81f2b17d 100644
--- a/gem/coffee.py
+++ b/gem/coffee.py
@@ -4,17 +4,17 @@
 This file is NOT for code generation as a COFFEE AST.
 """
 
-from collections import defaultdict
+from collections import Counter, defaultdict
 from itertools import chain, repeat
 import logging
 
 import numpy
 
-from gem.gem import ComponentTensor, Index, Indexed, IndexSum, one
+from gem.gem import ComponentTensor, Index, Indexed, IndexSum, Literal, Node, one
 from gem.node import MemoizerArg
 from gem.cost import has_arithmetic
 from gem.optimise import (filtered_replace_indices,
-                          make_sum, make_product)
+                          make_sum, make_product, traverse_sum)
 from gem.refactorise import Monomial, MonomialSum
 from gem.utils import groupby
 
@@ -201,6 +201,62 @@ def group_key(monomial):
     return new_monomials
 
 
+def _additive_map_key(
+        expression: Node) -> frozenset[tuple[Node, int]]:
+    """Represent an additive linear map independently of term order."""
+    return frozenset(Counter(traverse_sum(expression)).items())
+
+
+def _extract_repeated_linear_maps(
+        monomial_sum: MonomialSum,
+        linear_indices: tuple[Index, ...]) -> MonomialSum:
+    """Move uniform multiplicities from linear maps into scalar factors.
+
+    Parameters
+    ----------
+    monomial_sum
+        Sum-of-products representation of a multilinear expression.
+    linear_indices
+        Free indices identifying argument axes.
+
+    Returns
+    -------
+    MonomialSum
+        Representation with primitive linear maps and scalar multiplicities.
+
+    Notes
+    -----
+    A repeated additive map is a scalar multiple of the map formed from its
+    distinct summands.  Keeping that scalar in the monomial remainder leaves
+    the atomic factor as the finite element map that should be materialised.
+
+    """
+
+    linear_set = frozenset(linear_indices)
+    result = MonomialSum()
+    for monomial in monomial_sum:
+        atomics = []
+        factors = []
+        for atomic in monomial.atomics:
+            involved = linear_set.intersection(atomic.free_indices)
+            if len(involved) == 1:
+                summands = traverse_sum(atomic)
+                counts = Counter(summands)
+                multiplicities = set(counts.values())
+                if len(multiplicities) == 1:
+                    multiplicity, = multiplicities
+                    if multiplicity > 1:
+                        atomic = make_sum(counts)
+                        factors.append(Literal(float(multiplicity)))
+            atomics.append(atomic)
+        result.add(
+            monomial.sum_indices,
+            atomics,
+            make_product((*factors, monomial.rest)),
+        )
+    return result
+
+
 def _share_linear_maps(
         monomial_sum: MonomialSum,
         linear_indices: tuple[Index, ...]) -> MonomialSum:
@@ -228,6 +284,7 @@ def _share_linear_maps(
 
     """
     linear_indices = tuple(linear_indices)
+    monomial_sum = _extract_repeated_linear_maps(monomial_sum, linear_indices)
     linear_set = frozenset(linear_indices)
     canonical = {
         index.extent: Index(extent=index.extent)
@@ -235,6 +292,7 @@ def _share_linear_maps(
     }
     replacer = MemoizerArg(filtered_replace_indices)
     groups = defaultdict(list)
+    representatives = {}
     for monomial in monomial_sum:
         for atomic in monomial.atomics:
             involved = linear_set.intersection(atomic.free_indices)
@@ -242,10 +300,13 @@ def _share_linear_maps(
                 continue
             index, = involved
             normal = replacer(atomic, ((index, canonical[index.extent]),))
-            groups[normal].append((atomic, index))
+            key = _additive_map_key(normal)
+            groups[key].append((atomic, index))
+            representatives.setdefault(key, normal)
 
     replacements = {}
-    for normal, occurrences in groups.items():
+    for key, occurrences in groups.items():
+        normal = representatives[key]
         indices = {index for _, index in occurrences}
         if len(indices) < 2 or not has_arithmetic((normal,)):
             continue
diff --git a/test/gem/test_sum_factorise.py b/test/gem/test_sum_factorise.py
index 5dca72b9..3d524993 100644
--- a/test/gem/test_sum_factorise.py
+++ b/test/gem/test_sum_factorise.py
@@ -233,6 +233,34 @@ def test_preserved_linear_map_is_shared():
     assert numpy.allclose(result.arr, expected)
 
 
+def test_linear_map_sharing_normalises_order_and_scale() -> None:
+    i, j = gem.Index(extent=3), gem.Index(extent=3)
+    a_values = numpy.array([1.0, 2.0, 3.0])
+    b_values = numpy.array([4.0, 5.0, 6.0])
+    a = gem.Literal(a_values)
+    b = gem.Literal(b_values)
+    left_terms = (gem.Indexed(a, (i,)), gem.Indexed(b, (i,)))
+    right_terms = (gem.Indexed(b, (j,)), gem.Indexed(a, (j,)))
+    left_map = optimise.make_sum(left_terms * 2)
+    right_map = optimise.make_sum(right_terms * 2)
+    expression = gem.Product(left_map, right_map)
+
+    arguments = (i, j)
+    classifier = partial(_classify, frozenset(arguments))
+    preserved, = collect_monomials(
+        [expression], classifier, arguments)
+    optimised = optimise_monomial_sum(preserved, arguments)
+    tensors = [node for node in traversal((optimised,))
+               if isinstance(node, gem.ComponentTensor)]
+    assert len(tensors) == 1
+    body, = tensors[0].children
+    assert len(optimise.traverse_sum(body)) == 2
+
+    result, = evaluate([gem.ComponentTensor(optimised, arguments)])
+    expected = 4 * numpy.outer(a_values + b_values, a_values + b_values)
+    assert numpy.allclose(result.arr, expected)
+
+
 def test_expanded_and_preserved_agree():
     arguments, expanded, expected = monomial_sum(linear_indices=False)
     expression = optimise_monomial_sum(expanded, arguments)

From 205746ea95ec8a77af3295418dcf443d1401cbdd Mon Sep 17 00:00:00 2001
From: Pablo Brubeck <brubeck@protonmail.com>
Date: Fri, 21 Aug 2026 15:15:00 +0100
Subject: [PATCH 02/11] Factor reductions before indirect gathers

---
 gem/coffee.py       |  6 ++-
 gem/impero_utils.py |  7 +++-
 gem/optimise.py     | 89 +++++++++++++++++++++++++++++++++++++++++++++
 3 files changed, 98 insertions(+), 4 deletions(-)

diff --git a/gem/coffee.py b/gem/coffee.py
index 81f2b17d..33c40a63 100644
--- a/gem/coffee.py
+++ b/gem/coffee.py
@@ -13,7 +13,8 @@
 from gem.gem import ComponentTensor, Index, Indexed, IndexSum, Literal, Node, one
 from gem.node import MemoizerArg
 from gem.cost import has_arithmetic
-from gem.optimise import (filtered_replace_indices,
+from gem.optimise import (factorise_indirect_reductions,
+                          filtered_replace_indices,
                           make_sum, make_product, traverse_sum)
 from gem.refactorise import Monomial, MonomialSum
 from gem.utils import groupby
@@ -344,7 +345,8 @@ def optimise_monomial_sum(monomial_sum, linear_indices):
     new_monomials = []
     for _, monomials in groups:
         new_monomials.extend(optimise_monomials(monomials, linear_indices))
-    return monomial_sum_to_expression(new_monomials)
+    expression = monomial_sum_to_expression(new_monomials)
+    return factorise_indirect_reductions(expression)
 
 
 def optimise_monomials(monomials, linear_indices):
diff --git a/gem/impero_utils.py b/gem/impero_utils.py
index 31f9565b..2f15cd64 100644
--- a/gem/impero_utils.py
+++ b/gem/impero_utils.py
@@ -128,14 +128,17 @@ def inline_temporaries(expressions, ops):
     for op in ops:
         if isinstance(op, imp.Evaluate):
             expr = op.expression
-            if expr.shape == () and refcount[expr] == 1:
+            reduction_view = (isinstance(expr, gem.ComponentTensor)
+                              and isinstance(expr.children[0], gem.IndexSum))
+            if (expr.shape == () or reduction_view) and refcount[expr] == 1:
                 candidates.add(expr)
 
     # Prevent inlining that pulls expressions into inner loops
     for node in traversal(expressions):
         for child in node.children:
             if child in candidates and set(child.free_indices) < set(node.free_indices):
-                candidates.remove(child)
+                if not isinstance(child, gem.ComponentTensor):
+                    candidates.remove(child)
 
     # Filter out candidates
     return [op for op in ops if not (isinstance(op, imp.Evaluate) and op.expression in candidates)]
diff --git a/gem/optimise.py b/gem/optimise.py
index f009618d..32878f8d 100644
--- a/gem/optimise.py
+++ b/gem/optimise.py
@@ -1232,3 +1232,92 @@ def aggressive_unroll(expression):
     expression, = unroll_indexsum((expression,), predicate=lambda index: True)
     expression, = remove_componenttensors((expression,))
     return expression
+
+
+def factorise_indirect_reductions(expression: Node) -> Node:
+    """Factor a reduction through a repeated indirect table lookup.
+
+    Parameters
+    ----------
+    expression
+        Root of a scalar GEM expression.
+
+    Returns
+    -------
+    Node
+        Expression with profitable dense reductions evaluated before gathers.
+
+    Notes
+    -----
+    For an indirect index c(i), linearity lets the reduction over q be
+    evaluated for a new dense table-row index before gathering at c(i).
+    The rewrite is selected only when GEM's cost model predicts less
+    arithmetic, with storage and node count breaking ties.
+
+    """
+    def replace(node, target, replacement, cache):
+        key = id(node)
+        if key in cache:
+            return cache[key]
+        children = tuple(replace(child, target, replacement, cache)
+                         for child in node.children)
+        if isinstance(node, Indexed):
+            multiindex = tuple(
+                replacement if index == target else index
+                for index in node.multiindex)
+            result = Indexed(children[0], multiindex)
+        elif children == node.children:
+            result = node
+        else:
+            result = node.reconstruct(*children)
+        cache[key] = result
+        return result
+
+    def choose(node):
+        body, = node.children
+        candidates = OrderedDict()
+        for indexed in traversal((body,)):
+            if not isinstance(indexed, Indexed):
+                continue
+            aggregate, = indexed.children
+            for index, extent in zip(indexed.multiindex, aggregate.shape):
+                if isinstance(index, VariableIndex):
+                    sources = frozenset(index.expression.free_indices)
+                    if sources and sources.isdisjoint(node.multiindex):
+                        candidates.setdefault(index, (extent, sources))
+
+        best = node
+        best_cost = estimate_cost((node,))
+        for indirect, (extent, sources) in candidates.items():
+            latent = Index(extent=extent)
+            dense_body = replace(body, indirect, latent, {})
+            if sources.intersection(dense_body.free_indices):
+                continue
+            dense = ComponentTensor(
+                IndexSum(dense_body, node.multiindex), (latent,))
+            candidate = Indexed(dense, (indirect,))
+            if candidate.free_indices != node.free_indices:
+                continue
+            cost = estimate_cost((candidate,))
+            if cost < best_cost:
+                best = candidate
+                best_cost = cost
+        return best
+
+    cache = {}
+
+    def visit(node):
+        key = id(node)
+        if key in cache:
+            return cache[key]
+        children = tuple(visit(child) for child in node.children)
+        if children == node.children:
+            result = node
+        else:
+            result = node.reconstruct(*children)
+        if isinstance(result, IndexSum):
+            result = choose(result)
+        cache[key] = result
+        return result
+
+    return visit(expression)

From c339b9fbe3335d97da206d6f2b7d1a1ba6b25699 Mon Sep 17 00:00:00 2001
From: Pablo Brubeck <brubeck@protonmail.com>
Date: Sat, 22 Aug 2026 12:07:49 +0100
Subject: [PATCH 03/11] Factor mapped scalar coefficients

---
 finat/physically_mapped.py     |  4 +-
 gem/optimise.py                | 79 ++++++++++++++++++++++++++++++++++
 test/gem/test_sum_factorise.py | 30 +++++++++++++
 3 files changed, 112 insertions(+), 1 deletion(-)

diff --git a/finat/physically_mapped.py b/finat/physically_mapped.py
index da88e7d6..41ba4ea7 100644
--- a/finat/physically_mapped.py
+++ b/finat/physically_mapped.py
@@ -4,6 +4,7 @@
 
 import gem
 import numpy
+from gem.optimise import factorise_scalar_sums
 
 from finat.citations import cite
 
@@ -66,7 +67,8 @@ def __init__(
         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)
+            data[index, :len(row)] = tuple(
+                factorise_scalar_sums(gem.as_gem(value)) for _, value in row)
         self._width = width
         self._columns = gem.Literal(columns, dtype=gem.uint_type)
         values = []
diff --git a/gem/optimise.py b/gem/optimise.py
index 32878f8d..935facb7 100644
--- a/gem/optimise.py
+++ b/gem/optimise.py
@@ -1234,6 +1234,85 @@ def aggressive_unroll(expression):
     return expression
 
 
+def factorise_scalar_sums(expression: Node) -> Node:
+    """Factor common products from scalar sums when this lowers GEM cost.
+
+    Parameters
+    ----------
+    expression
+        Root of a GEM expression.
+
+    Returns
+    -------
+    Node
+        Expression with profitable common product factors extracted.
+
+    Notes
+    -----
+    Scalar geometry and basis-transformation expressions are simplified below
+    the indexed contraction structure.  Contractions are indivisible factors:
+    their bound indices cannot move through an enclosing sum.  Sums carrying
+    free indices are left to the contraction planner, whose cost model includes
+    their iteration domains.
+
+    """
+    def choose(node):
+        if node.free_indices:
+            return node
+        summands = traverse_sum(node)
+        if len(summands) < 2:
+            return node
+
+        factorisations = []
+        for summand in summands:
+            _, factors = traverse_product(
+                summand,
+                stop_at=lambda factor: isinstance(factor, IndexSum),
+            )
+            factorisations.append(factors)
+
+        common = Counter(factorisations[0])
+        for factors in factorisations[1:]:
+            common &= Counter(factors)
+        if not common:
+            return node
+
+        common_factors = list(common.elements())
+        remainders = []
+        for factors in factorisations:
+            remaining_common = common.copy()
+            remaining = []
+            for factor in factors:
+                if remaining_common[factor]:
+                    remaining_common[factor] -= 1
+                else:
+                    remaining.append(factor)
+            remainders.append(make_product(remaining))
+
+        candidate = make_product(
+            (*common_factors, make_sum(remainders)))
+        if candidate.free_indices != node.free_indices:
+            return node
+        if estimate_cost((candidate,)) < estimate_cost((node,)):
+            return candidate
+        return node
+
+    cache = {}
+
+    def visit(node):
+        key = id(node)
+        if key in cache:
+            return cache[key]
+        children = tuple(visit(child) for child in node.children)
+        result = node if children == node.children else node.reconstruct(*children)
+        if isinstance(result, Sum):
+            result = choose(result)
+        cache[key] = result
+        return result
+
+    return visit(expression)
+
+
 def factorise_indirect_reductions(expression: Node) -> Node:
     """Factor a reduction through a repeated indirect table lookup.
 
diff --git a/test/gem/test_sum_factorise.py b/test/gem/test_sum_factorise.py
index 3d524993..9664c761 100644
--- a/test/gem/test_sum_factorise.py
+++ b/test/gem/test_sum_factorise.py
@@ -300,3 +300,33 @@ def test_contraction_counts_index_absent_from_factors() -> None:
     result, = evaluate([expression])
 
     assert result.arr == 2 * numpy.arange(3.0).sum()
+
+
+def test_common_factor_is_extracted_from_scalar_sum() -> None:
+    common = gem.Variable("common", ())
+    left = gem.Variable("left", ())
+    right = gem.Variable("right", ())
+    expression = gem.Sum(
+        gem.Product(common, left), gem.Product(common, right))
+
+    expression = optimise.factorise_scalar_sums(expression)
+
+    _, factors = optimise.traverse_product(expression)
+    assert common in factors
+    result, = evaluate([expression], {
+        common: numpy.asarray(2.0),
+        left: numpy.asarray(3.0),
+        right: numpy.asarray(4.0),
+    })
+    assert result.arr == 14
+
+
+def test_scalar_factorisation_leaves_indexed_sum_to_planner() -> None:
+    i = gem.Index(extent=3)
+    common = gem.Variable("common", ())
+    left = gem.Indexed(gem.Literal(numpy.arange(3.0)), (i,))
+    right = gem.Indexed(gem.Literal(numpy.arange(3.0, 6.0)), (i,))
+    expression = gem.Sum(
+        gem.Product(common, left), gem.Product(common, right))
+
+    assert optimise.factorise_scalar_sums(expression) is expression

From 7541d617a6c5a26087f97c59e6d9f0d991f007b2 Mon Sep 17 00:00:00 2001
From: Pablo Brubeck <brubeck@protonmail.com>
Date: Sat, 22 Aug 2026 16:55:40 +0100
Subject: [PATCH 04/11] Place indirect reductions once per assignment

optimise_monomial_sum runs inside the recursive sum_factorise, so calling
the reduction placement there walked the same subtrees once per recursion
level: 754 calls over 75,497 nodes for a Johnson-Mercier 3D matrix, which
the pass never improves.  Leave the traversal to the caller, memoise it on
the node rather than its id, and skip costing an IndexSum that offers no
indirect gather to factor.

Co-Authored-By: Claude Opus 5 <noreply@anthropic.com>
---
 gem/coffee.py   |  6 ++----
 gem/optimise.py | 50 +++++++++++++++++--------------------------------
 2 files changed, 19 insertions(+), 37 deletions(-)

diff --git a/gem/coffee.py b/gem/coffee.py
index 33c40a63..81f2b17d 100644
--- a/gem/coffee.py
+++ b/gem/coffee.py
@@ -13,8 +13,7 @@
 from gem.gem import ComponentTensor, Index, Indexed, IndexSum, Literal, Node, one
 from gem.node import MemoizerArg
 from gem.cost import has_arithmetic
-from gem.optimise import (factorise_indirect_reductions,
-                          filtered_replace_indices,
+from gem.optimise import (filtered_replace_indices,
                           make_sum, make_product, traverse_sum)
 from gem.refactorise import Monomial, MonomialSum
 from gem.utils import groupby
@@ -345,8 +344,7 @@ def optimise_monomial_sum(monomial_sum, linear_indices):
     new_monomials = []
     for _, monomials in groups:
         new_monomials.extend(optimise_monomials(monomials, linear_indices))
-    expression = monomial_sum_to_expression(new_monomials)
-    return factorise_indirect_reductions(expression)
+    return monomial_sum_to_expression(new_monomials)
 
 
 def optimise_monomials(monomials, linear_indices):
diff --git a/gem/optimise.py b/gem/optimise.py
index 935facb7..10b54d8a 100644
--- a/gem/optimise.py
+++ b/gem/optimise.py
@@ -1334,23 +1334,14 @@ def factorise_indirect_reductions(expression: Node) -> Node:
     arithmetic, with storage and node count breaking ties.
 
     """
-    def replace(node, target, replacement, cache):
-        key = id(node)
-        if key in cache:
-            return cache[key]
-        children = tuple(replace(child, target, replacement, cache)
-                         for child in node.children)
+    def replace(node, self, substitution):
+        target, replacement = substitution
         if isinstance(node, Indexed):
-            multiindex = tuple(
-                replacement if index == target else index
-                for index in node.multiindex)
-            result = Indexed(children[0], multiindex)
-        elif children == node.children:
-            result = node
-        else:
-            result = node.reconstruct(*children)
-        cache[key] = result
-        return result
+            child, = node.children
+            multiindex = tuple(replacement if index == target else index
+                               for index in node.multiindex)
+            return Indexed(self(child, substitution), multiindex)
+        return reuse_if_untouched_arg(node, self, substitution)
 
     def choose(node):
         body, = node.children
@@ -1365,11 +1356,14 @@ def choose(node):
                     if sources and sources.isdisjoint(node.multiindex):
                         candidates.setdefault(index, (extent, sources))
 
+        if not candidates:
+            return node
+
         best = node
         best_cost = estimate_cost((node,))
         for indirect, (extent, sources) in candidates.items():
             latent = Index(extent=extent)
-            dense_body = replace(body, indirect, latent, {})
+            dense_body = MemoizerArg(replace)(body, (indirect, latent))
             if sources.intersection(dense_body.free_indices):
                 continue
             dense = ComponentTensor(
@@ -1383,20 +1377,10 @@ def choose(node):
                 best_cost = cost
         return best
 
-    cache = {}
-
-    def visit(node):
-        key = id(node)
-        if key in cache:
-            return cache[key]
-        children = tuple(visit(child) for child in node.children)
-        if children == node.children:
-            result = node
-        else:
-            result = node.reconstruct(*children)
-        if isinstance(result, IndexSum):
-            result = choose(result)
-        cache[key] = result
-        return result
+    def visit(node, self):
+        node = reuse_if_untouched(node, self)
+        if isinstance(node, IndexSum):
+            node = choose(node)
+        return node
 
-    return visit(expression)
+    return Memoizer(visit)(expression)

From dd4c4d76f7be075747f65e1dde31d5c23fe683e5 Mon Sep 17 00:00:00 2001
From: Pablo Brubeck <brubeck@protonmail.com>
Date: Mon, 24 Aug 2026 17:16:27 +0100
Subject: [PATCH 05/11] tabulate_indirect_contractions

---
 gem/optimise.py | 109 +++++++++++++++++++++++++-----------------------
 1 file changed, 57 insertions(+), 52 deletions(-)

diff --git a/gem/optimise.py b/gem/optimise.py
index 10b54d8a..980a8d65 100644
--- a/gem/optimise.py
+++ b/gem/optimise.py
@@ -1313,74 +1313,79 @@ def visit(node):
     return visit(expression)
 
 
-def factorise_indirect_reductions(expression: Node) -> Node:
-    """Factor a reduction through a repeated indirect table lookup.
+def _indirect_gathers(expression: Node) -> OrderedDict:
+    """Find the maps that an expression reads tables through.
 
-    Parameters
-    ----------
-    expression
-        Root of a scalar GEM expression.
+    :arg expression: the root of a scalar GEM expression
+    :returns: an ordered mapping from each :class:`~.VariableIndex` that
+              indexes a table to the number of rows it selects from
+    """
+    gathers = OrderedDict()
+    for node in traversal((expression,)):
+        if isinstance(node, Indexed):
+            aggregate, = node.children
+            for index, extent in zip(node.multiindex, aggregate.shape):
+                if isinstance(index, VariableIndex) and index.expression.free_indices:
+                    gathers.setdefault(index, extent)
+    return gathers
 
-    Returns
-    -------
-    Node
-        Expression with profitable dense reductions evaluated before gathers.
 
-    Notes
-    -----
-    For an indirect index c(i), linearity lets the reduction over q be
-    evaluated for a new dense table-row index before gathering at c(i).
-    The rewrite is selected only when GEM's cost model predicts less
-    arithmetic, with storage and node count breaking ties.
+def tabulate_indirect_contractions(expression: Node) -> Node:
+    """Evaluate a contraction one time for each row of its table.
+
+    An expression can read a table through a map.  Each argument of the map
+    selects one row, and two arguments can select the same row.  A
+    contraction that reads the table through the map thus repeats work.
+
+    This function evaluates the contraction for each row, then reads those
+    results through the map.  The rewrite is correct only if the map does not
+    change with the contracted indices, and no other part of the contraction
+    uses the argument.  It is faster only if the map has more arguments than
+    the table has rows.
 
+    :arg expression: the root of a scalar GEM expression
+    :returns: the expression with each such contraction evaluated one time
+              for each row
     """
-    def replace(node, self, substitution):
+    gathers = _indirect_gathers(expression)
+    if not gathers:
+        return expression
+
+    def rename(node, self, substitution):
         target, replacement = substitution
         if isinstance(node, Indexed):
-            child, = node.children
+            aggregate, = node.children
             multiindex = tuple(replacement if index == target else index
                                for index in node.multiindex)
-            return Indexed(self(child, substitution), multiindex)
+            return Indexed(self(aggregate, substitution), multiindex)
         return reuse_if_untouched_arg(node, self, substitution)
 
-    def choose(node):
+    def hoist(node):
         body, = node.children
-        candidates = OrderedDict()
-        for indexed in traversal((body,)):
-            if not isinstance(indexed, Indexed):
-                continue
-            aggregate, = indexed.children
-            for index, extent in zip(indexed.multiindex, aggregate.shape):
-                if isinstance(index, VariableIndex):
-                    sources = frozenset(index.expression.free_indices)
-                    if sources and sources.isdisjoint(node.multiindex):
-                        candidates.setdefault(index, (extent, sources))
-
-        if not candidates:
-            return node
-
-        best = node
-        best_cost = estimate_cost((node,))
-        for indirect, (extent, sources) in candidates.items():
-            latent = Index(extent=extent)
-            dense_body = MemoizerArg(replace)(body, (indirect, latent))
-            if sources.intersection(dense_body.free_indices):
-                continue
-            dense = ComponentTensor(
-                IndexSum(dense_body, node.multiindex), (latent,))
-            candidate = Indexed(dense, (indirect,))
-            if candidate.free_indices != node.free_indices:
-                continue
-            cost = estimate_cost((candidate,))
-            if cost < best_cost:
-                best = candidate
-                best_cost = cost
-        return best
+        free = frozenset(body.free_indices)
+        contracted = frozenset(node.multiindex)
+
+        candidates = []
+        for gather, nrows in gathers.items():
+            arguments = frozenset(gather.expression.free_indices)
+            if arguments <= free and arguments.isdisjoint(contracted):
+                saving = _iteration_count(arguments) - nrows
+                if saving > 0:
+                    candidates.append((saving, gather, arguments, nrows))
+
+        for _, gather, arguments, nrows in sorted(candidates, key=lambda c: -c[0]):
+            row = Index(extent=nrows)
+            per_row = MemoizerArg(rename)(body, (gather, row))
+            # The body must reach the arguments only through this gather.
+            if arguments.isdisjoint(per_row.free_indices):
+                table = ComponentTensor(IndexSum(per_row, node.multiindex), (row,))
+                return Indexed(table, (gather,))
+        return node
 
     def visit(node, self):
         node = reuse_if_untouched(node, self)
         if isinstance(node, IndexSum):
-            node = choose(node)
+            node = hoist(node)
         return node
 
     return Memoizer(visit)(expression)

From f1dd25a13cfbaafa08c4f86341c374741f4d5596 Mon Sep 17 00:00:00 2001
From: Pablo Brubeck <brubeck@protonmail.com>
Date: Mon, 24 Aug 2026 18:00:19 +0100
Subject: [PATCH 06/11] Match the docstring style of the files they are in

Co-Authored-By: Claude Opus 5 <noreply@anthropic.com>
---
 finat/physically_mapped.py |  47 +++++-----------
 gem/coffee.py              |  48 +++++-----------
 gem/optimise.py            | 112 +++++++++----------------------------
 gem/refactorise.py         |  78 +++++++++-----------------
 4 files changed, 81 insertions(+), 204 deletions(-)

diff --git a/finat/physically_mapped.py b/finat/physically_mapped.py
index 41ba4ea7..8027732b 100644
--- a/finat/physically_mapped.py
+++ b/finat/physically_mapped.py
@@ -22,25 +22,17 @@ def dual_transformation(self, Q, coordinate_mapping=None):
 
 
 class MappedTabulation(Mapping):
-    """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.
+    """A lazy tabulation dict that applies the basis transformation only
+    on the requested derivatives.
 
+    Rows are padded to a common number of entries, so that a loop over the
+    basis index has an affine iteration domain.  Constant tables select the
+    reference column and one of the distinct symbolic coefficients, which
+    shares equal entries without materialising the 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__(
@@ -99,20 +91,11 @@ def _row_index(self) -> gem.Index:
     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.
-
+        :arg r: index over the rows retained by the element
+        :arg a: index over the reference basis
+        :returns: 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
         """
         k = self._row_index
         entry = gem.Indexed(
diff --git a/gem/coffee.py b/gem/coffee.py
index 81f2b17d..418b4926 100644
--- a/gem/coffee.py
+++ b/gem/coffee.py
@@ -212,26 +212,16 @@ def _extract_repeated_linear_maps(
         linear_indices: tuple[Index, ...]) -> MonomialSum:
     """Move uniform multiplicities from linear maps into scalar factors.
 
-    Parameters
-    ----------
-    monomial_sum
-        Sum-of-products representation of a multilinear expression.
-    linear_indices
-        Free indices identifying argument axes.
-
-    Returns
-    -------
-    MonomialSum
-        Representation with primitive linear maps and scalar multiplicities.
-
-    Notes
-    -----
     A repeated additive map is a scalar multiple of the map formed from its
     distinct summands.  Keeping that scalar in the monomial remainder leaves
     the atomic factor as the finite element map that should be materialised.
 
+    :arg monomial_sum: sum-of-products representation of a multilinear
+                       expression
+    :arg linear_indices: free indices identifying argument axes
+    :returns: representation with primitive linear maps and scalar
+              multiplicities
     """
-
     linear_set = frozenset(linear_indices)
     result = MonomialSum()
     for monomial in monomial_sum:
@@ -262,26 +252,16 @@ def _share_linear_maps(
         linear_indices: tuple[Index, ...]) -> MonomialSum:
     """Share isomorphic maps of distinct multilinear axes.
 
-    Parameters
-    ----------
-    monomial_sum
-        Sum-of-products representation of a multilinear expression.
-    linear_indices
-        Free indices identifying argument axes.
-
-    Returns
-    -------
-    MonomialSum
-        Representation whose repeated linear maps access one tensor.
-
-    Notes
-    -----
     Test and trial axes use distinct indices even when they apply the same
-    finite element map.  Renaming each axis to a canonical index exposes
-    that isomorphism without inspecting the element family.  Materialising
-    the canonical map is generalised code motion: the basis transformation
-    is evaluated once and both axes index its result.
-
+    finite element map.  Renaming each axis to a canonical index exposes that
+    isomorphism without inspecting the element family.  Materialising the
+    canonical map is generalised code motion: the basis transformation is
+    evaluated once and both axes index its result.
+
+    :arg monomial_sum: sum-of-products representation of a multilinear
+                       expression
+    :arg linear_indices: free indices identifying argument axes
+    :returns: representation whose repeated linear maps access one tensor
     """
     linear_indices = tuple(linear_indices)
     monomial_sum = _extract_repeated_linear_maps(monomial_sum, linear_indices)
diff --git a/gem/optimise.py b/gem/optimise.py
index 980a8d65..63a65fef 100644
--- a/gem/optimise.py
+++ b/gem/optimise.py
@@ -9,7 +9,7 @@
 
 import numpy
 
-from gem.cost import index_space_literal
+from gem.cost import estimate_cost, index_space_literal, iteration_count
 from gem.utils import groupby
 from gem.node import (Memoizer, MemoizerArg, reuse_if_untouched,
                       reuse_if_untouched_arg, traversal, traversal_children)
@@ -767,24 +767,11 @@ def traverse_sum(expression, stop_at=None):
 def distribute_sum(expr: Node, predicate: Callable[[Node], bool]) -> list[Node]:
     """Distribute selected sums through products and contractions.
 
-    Parameters
-    ----------
-    expr
-        GEM expression to distribute.
-    predicate
-        Predicate selecting the operations to distribute.
-
-    Returns
-    -------
-    list of Node
-        Additive terms after distribution.
-
-    Notes
-    -----
-    Memoization uses object identity.  Structurally equal GEM nodes can have
-    deep expression trees, while distribution only needs to reuse actual DAG
-    nodes.
+    Memoisation is by object identity, which reuses each node of the DAG once.
 
+    :arg expr: GEM expression to distribute
+    :arg predicate: predicate selecting the operations to distribute
+    :returns: the additive terms after distribution
     """
     results = {}
     active = {}
@@ -826,18 +813,9 @@ def distribute_sum(expr: Node, predicate: Callable[[Node], bool]) -> list[Node]:
 def _is_linear_map(node: Node, linear_indices: frozenset) -> bool:
     """Is a node a linear map into one multilinear axis?
 
-    Parameters
-    ----------
-    node
-        GEM expression node.
-    linear_indices
-        Free indices identifying the multilinear axes.
-
-    Returns
-    -------
-    bool
-        Whether the node is a sum over exactly one such axis.
-
+    :arg node: GEM expression node
+    :arg linear_indices: free indices identifying the multilinear axes
+    :returns: whether the node is a sum over exactly one such axis
     """
     return (isinstance(node, Sum)
             and len(linear_indices.intersection(node.free_indices)) == 1)
@@ -848,23 +826,12 @@ def has_linear_maps(
         linear_indices: Iterable[Index]) -> bool:
     """Does a GEM DAG contain a finite element linear map?
 
-    Parameters
-    ----------
-    expressions
-        Roots of a multilinear GEM expression DAG.
-    linear_indices
-        Free indices identifying the multilinear axes.
-
-    Returns
-    -------
-    bool
-        Whether preserving one-axis sums can change the factorisation.
-
-    Notes
-    -----
-    Answering this costs one traversal, where building the preserved
-    factorisation to compare it costs a whole pass of monomial collection.
+    One traversal answers this, so a caller can gate a second factorisation
+    on it.
 
+    :arg expressions: roots of a multilinear GEM expression DAG
+    :arg linear_indices: free indices identifying the multilinear axes
+    :returns: whether preserving one-axis sums can change the factorisation
     """
     linear_indices = frozenset(linear_indices)
     return any(_is_linear_map(node, linear_indices)
@@ -878,22 +845,13 @@ def preserve_linear_maps(
     """Expose multilinear terms and retain each one-axis linear map.
 
     A sum that depends on one linear index represents a linear map into an
-    argument tabulation. A sum that depends on several linear indices
-    separates multilinear form terms. This function distributes the latter
+    argument tabulation.  A sum that depends on several linear indices
+    separates multilinear form terms.  This function distributes the latter
     sums and returns the former sums as factors.
 
-    Parameters
-    ----------
-    expression
-        Multilinear GEM expression.
-    linear_indices
-        Free indices identifying the linear axes.
-
-    Returns
-    -------
-    tuple
-        Additive terms and the linear-map factors that they contain.
-
+    :arg expression: multilinear GEM expression
+    :arg linear_indices: free indices identifying the linear axes
+    :returns: the additive terms, and the linear-map factors they contain
     """
     linear_indices = frozenset(linear_indices)
 
@@ -958,17 +916,9 @@ def _delta_axes(node: Node, self: Memoizer) -> frozenset:
 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.
-
+    :arg index: index to inspect
+    :returns: ``(array, indices)`` when the index is a lookup into a
+              :class:`~.Literal` with a plain multiindex, otherwise ``None``
     """
     if not isinstance(index, VariableIndex):
         return None
@@ -1237,24 +1187,14 @@ def aggressive_unroll(expression):
 def factorise_scalar_sums(expression: Node) -> Node:
     """Factor common products from scalar sums when this lowers GEM cost.
 
-    Parameters
-    ----------
-    expression
-        Root of a GEM expression.
-
-    Returns
-    -------
-    Node
-        Expression with profitable common product factors extracted.
-
-    Notes
-    -----
     Scalar geometry and basis-transformation expressions are simplified below
     the indexed contraction structure.  Contractions are indivisible factors:
     their bound indices cannot move through an enclosing sum.  Sums carrying
-    free indices are left to the contraction planner, whose cost model includes
-    their iteration domains.
+    free indices are left to the contraction planner, whose cost model
+    includes their iteration domains.
 
+    :arg expression: root of a GEM expression
+    :returns: the expression with profitable common product factors extracted
     """
     def choose(node):
         if node.free_indices:
@@ -1369,7 +1309,7 @@ def hoist(node):
         for gather, nrows in gathers.items():
             arguments = frozenset(gather.expression.free_indices)
             if arguments <= free and arguments.isdisjoint(contracted):
-                saving = _iteration_count(arguments) - nrows
+                saving = iteration_count(arguments) - nrows
                 if saving > 0:
                     candidates.append((saving, gather, arguments, nrows))
 
diff --git a/gem/refactorise.py b/gem/refactorise.py
index b0b8a914..da9f1166 100644
--- a/gem/refactorise.py
+++ b/gem/refactorise.py
@@ -272,30 +272,15 @@ def _collect_monomial_sums(
         classifier: Callable[[Node], str]) -> list[MonomialSum]:
     """Collect monomial sums using the supplied node classifier.
 
-    Parameters
-    ----------
-    expressions : iterable of Node
-        GEM expressions to refactorise.
-    classifier : callable
-        Function labelling each node as ``ATOMIC``, ``COMPOUND``, or
-        ``OTHER``.
-
-    Returns
-    -------
-    list of MonomialSum
-        Polynomial representations of the expressions.
-
-    Raises
-    ------
-    FactorisationError
-        If a compound expression cannot be expanded.
-
-    Notes
-    -----
     ``expressions`` must already have had its ComponentTensors removed, so
-    that a caller identifying nodes to classify sees the nodes this
-    collector will visit.
-
+    that a caller identifying nodes to classify sees the nodes this collector
+    will visit.
+
+    :arg expressions: GEM expressions to refactorise
+    :arg classifier: a function labelling each node as ``ATOMIC``,
+                     ``COMPOUND``, or ``OTHER``
+    :returns: list of :py:class:`MonomialSum`s
+    :raises FactorisationError: a compound expression cannot be expanded.
     """
 
     # Get ListTensors out of the way
@@ -322,36 +307,25 @@ def collect_monomials(
         expressions: Iterable[Node],
         classifier: Callable[[Node], str],
         linear_indices: Iterable[Index] = ()) -> list[MonomialSum]:
-    """Collect structure-preserving sum-of-products representations.
-
-    Parameters
-    ----------
-    expressions
-        GEM expressions to refactorise.
-    classifier
-        Function that labels GEM nodes for polynomial collection.
-    linear_indices
-        Free indices identifying the multilinear axes.  Sums depending on
-        exactly one such axis represent linear maps and remain atomic.
-
-    Returns
-    -------
-    list of MonomialSum
-        One polynomial representation for each input expression.
-
-    Notes
-    -----
-    A one-axis sum is a finite element linear operand: examples include a
-    sparse basis transformation and a tensor-product tabulation factor.
-    Preserving it keeps domain-specific basis structure available for
-    contraction optimisation.  This does not make the map opaque: its GEM
-    expression remains available for scalar simplification and code motion.
-
-    Sums involving several linear axes separate form monomials and are
-    distributed.  COFFEE subsequently chooses scalar factorisations across
-    the resulting operands.  Keeping these two stages distinct avoids a
-    Cartesian expansion of basis-map entries before loop placement.
+    """Refactorises expressions into a sum-of-products form, using
+    distributivity rules (i.e. a*(b + c) -> a*b + a*c).  Expansion
+    proceeds until all "compound" expressions are broken up.
+
+    A sum over exactly one linear index is a finite element linear map, such
+    as a sparse basis transformation or a tensor-product tabulation factor,
+    and stays atomic.  A sum over several linear indices separates form
+    monomials, and is distributed.
 
+    :arg expressions: GEM expressions to refactorise
+    :arg classifier: a function that can classify any GEM expression
+                     as ``ATOMIC``, ``COMPOUND``, or ``OTHER``.  This
+                     classification drives the factorisation.
+    :arg linear_indices: free indices identifying the multilinear axes
+
+    :returns: list of :py:class:`MonomialSum`s
+
+    :raises FactorisationError: Failed to break up some "compound"
+                                expressions with expansion.
     """
     # Remove ComponentTensors here, not in the collector.  Removing them
     # rebuilds nodes, and the maps found below must be the very nodes that

From ac5295d339a51d1ebbe211d27a5238046687e023 Mon Sep 17 00:00:00 2001
From: Pablo Brubeck <brubeck@protonmail.com>
Date: Mon, 24 Aug 2026 18:00:19 +0100
Subject: [PATCH 07/11] Memoise factorise_scalar_sums on the node

Sharing the result between structurally equal subexpressions, as the
other rewrites in this module do.

Co-Authored-By: Claude Opus 5 <noreply@anthropic.com>
---
 gem/optimise.py | 19 ++++++-------------
 1 file changed, 6 insertions(+), 13 deletions(-)

diff --git a/gem/optimise.py b/gem/optimise.py
index 63a65fef..a2133f5e 100644
--- a/gem/optimise.py
+++ b/gem/optimise.py
@@ -1237,20 +1237,13 @@ def choose(node):
             return candidate
         return node
 
-    cache = {}
+    def visit(node, self):
+        node = reuse_if_untouched(node, self)
+        if isinstance(node, Sum):
+            node = choose(node)
+        return node
 
-    def visit(node):
-        key = id(node)
-        if key in cache:
-            return cache[key]
-        children = tuple(visit(child) for child in node.children)
-        result = node if children == node.children else node.reconstruct(*children)
-        if isinstance(result, Sum):
-            result = choose(result)
-        cache[key] = result
-        return result
-
-    return visit(expression)
+    return Memoizer(visit)(expression)
 
 
 def _indirect_gathers(expression: Node) -> OrderedDict:

From 7ab701a9cd074702e1f09c32f87c5bb0518d429f Mon Sep 17 00:00:00 2001
From: Pablo Brubeck <brubeck@protonmail.com>
Date: Mon, 24 Aug 2026 18:00:19 +0100
Subject: [PATCH 08/11] Break contraction plan ties on the caller's index order

The tie in _plan_contraction ranks indices by their position in
sum_indices, which spectral.py fixes to quadrature source order and
_independent_contractions preserves into each subproblem.  A plan
therefore depends only on the contraction it is planning.

Co-Authored-By: Claude Opus 5 <noreply@anthropic.com>
---
 gem/optimise.py                |  5 ++++-
 test/gem/test_sum_factorise.py | 24 +++++++++++++++++++++++-
 2 files changed, 27 insertions(+), 2 deletions(-)

diff --git a/gem/optimise.py b/gem/optimise.py
index a2133f5e..8f2e6914 100644
--- a/gem/optimise.py
+++ b/gem/optimise.py
@@ -487,9 +487,12 @@ def size(indices):
         reducible[subset] = frozenset(index for index in sum_indices
                                       if not incidence[index] & ~subset)
 
+    # Rank of each contraction index in the caller's ordering, to break ties.
+    position = {index: n for n, index in enumerate(sum_indices)}
+
     def order(indices):
         """Sum out the widest index first, breaking ties reproducibly."""
-        return tuple(sorted(indices, key=lambda i: (-extents[i], i.count)))
+        return tuple(sorted(indices, key=lambda i: (-extents[i], position[i])))
 
     def reduce_indices(expression, free, indices):
         """Sum out indices, largest extent first, costing each reduction."""
diff --git a/test/gem/test_sum_factorise.py b/test/gem/test_sum_factorise.py
index 9664c761..33ac1991 100644
--- a/test/gem/test_sum_factorise.py
+++ b/test/gem/test_sum_factorise.py
@@ -7,7 +7,7 @@
 import gem
 from gem.interpreter import evaluate
 from gem import cost, optimise
-from gem.node import traversal
+from gem.node import post_traversal, traversal
 from gem.optimise import sum_factorise
 from gem.coffee import optimise_monomial_sum
 from gem.refactorise import (ATOMIC, COMPOUND, OTHER,
@@ -330,3 +330,25 @@ def test_scalar_factorisation_leaves_indexed_sum_to_planner() -> None:
         gem.Product(common, left), gem.Product(common, right))
 
     assert optimise.factorise_scalar_sums(expression) is expression
+
+
+def test_contraction_plan_follows_the_callers_index_order() -> None:
+    """A plan depends only on the caller's ordering of the sum indices."""
+    def plan(reverse_creation):
+        if reverse_creation:
+            r = gem.Index(extent=4)
+            q = gem.Index(extent=4)
+        else:
+            q = gem.Index(extent=4)
+            r = gem.Index(extent=4)
+        # Both factors carry both indices, so the two reduce at the same
+        # subset and the tie decides their order within one IndexSum.
+        factors = [gem.Indexed(gem.Variable(name, (4, 4)), (q, r))
+                   for name in ("A", "B")]
+        expression = optimise.sum_factorise([q, r], factors)
+        role = {q: "q", r: "r"}
+        return [tuple(role[i] for i in node.multiindex)
+                for node in post_traversal((expression,))
+                if isinstance(node, gem.IndexSum)]
+
+    assert plan(False) == plan(True) == [("q", "r")]

From d26b09ae12558dc2d2e82d8d5b509885118071ba Mon Sep 17 00:00:00 2001
From: Pablo Brubeck <brubeck@protonmail.com>
Date: Sat, 29 Aug 2026 10:37:29 +0100
Subject: [PATCH 09/11] Rename the indices an indirect gather hides

Indexed(ComponentTensor(Indexed(C, kk), jj), ii) rewrites kk in place, which
only reaches an index that appears in kk as itself.  An index reached through
a VariableIndex is a free index of the lookup, not a key of kk, so the rewrite
left it behind while the caller went on believing it had been renamed.

That is how make_renamer separates two sums that bind the same index, so the
two sums stayed joined:

    (sum_k a_k) * (sum_k b_k)  ==>  sum_k a_k b_k

Nothing exercised this until a basis transformation became a gather through a
lookup table.  Assembling a diagonal indexes both arguments with the test
index, so a mapped tabulation meets itself in a product, and the shared row
index turned the square of a sum into a sum of squares.  The matrix-free
diagonal of an HCT-red biharmonic form came out 5.4 times too large, and the
multigrid built on it took 20 iterations where 16 were expected.

Leave the pattern alone when jj is hidden in a VariableIndex, so that
replace_indices, which does substitute inside the lookup, handles it.

Co-Authored-By: Claude Opus 5 <noreply@anthropic.com>
Claude-Session: https://claude.ai/code/session_013NoTXWyj2fVdJTHnMDFB4k
---
 gem/gem.py                |  9 ++++++++-
 test/gem/test_simplify.py | 38 ++++++++++++++++++++++++++++++++++++++
 2 files changed, 46 insertions(+), 1 deletion(-)

diff --git a/gem/gem.py b/gem/gem.py
index d3d8c3a0..0953d697 100644
--- a/gem/gem.py
+++ b/gem/gem.py
@@ -727,7 +727,14 @@ def __new__(cls, aggregate, multiindex):
                 C, = B.children
                 kk = B.multiindex
                 ff = C.free_indices
-                if not any((j in ff) for j in jj):
+                # An index reached only through a VariableIndex is not a key
+                # of kk, so the rewrite below would drop its replacement and
+                # leave the old index behind, bound by nothing.  Those belong
+                # to replace_indices, which does substitute inside the lookup.
+                hh = frozenset(chain.from_iterable(
+                    k.expression.free_indices
+                    for k in kk if isinstance(k, VariableIndex)))
+                if not any((j in ff or j in hh) for j in jj):
                     # Only replace indices that are not present in C
                     rep = dict(zip(jj, ii))
                     ll = tuple(rep.get(k, k) for k in kk)
diff --git a/test/gem/test_simplify.py b/test/gem/test_simplify.py
index 94aeac3d..4343b1a4 100644
--- a/test/gem/test_simplify.py
+++ b/test/gem/test_simplify.py
@@ -100,3 +100,41 @@ def test_flatten_indexsum(A):
     result = gem.IndexSum(gem.IndexSum(Aij, (i,)), (j,))
     expected = gem.IndexSum(Aij, (i, j))
     assert result == expected
+
+
+def test_rename_index_under_variable_index():
+    """Renaming a bound index must reach the lookup of an indirect gather."""
+    values = gem.Literal(numpy.array([10.0, 20.0, 30.0]))
+    lookup = gem.Literal(numpy.array([2, 0, 1], dtype=gem.uint_type),
+                         dtype=gem.uint_type)
+
+    k, kp = gem.indices(2)
+    gather = gem.Indexed(
+        values, (gem.VariableIndex(gem.Indexed(lookup, (k,))),))
+    assert k in gather.free_indices
+
+    # This is how make_renamer separates two sums that bind the same index.
+    renamed = gem.Indexed(gem.ComponentTensor(gather, (k,)), (kp,))
+    assert renamed.free_indices == (kp,)
+
+
+def test_product_of_sums_over_one_index():
+    """Two sums that bind the same index expand to a double sum, not one."""
+    from gem.interpreter import evaluate
+    from gem.optimise import make_rename_map, make_renamer
+
+    values = gem.Literal(numpy.array([1.0, 2.0, 3.0]))
+    lookup = gem.Literal(numpy.array([2, 0, 1], dtype=gem.uint_type),
+                         dtype=gem.uint_type)
+    k, = gem.indices(1)
+    gather = gem.Indexed(
+        values, (gem.VariableIndex(gem.Indexed(lookup, (k,))),))
+
+    renamer = make_renamer(make_rename_map())
+    (k0,), first = renamer((k,))
+    (k1,), second = renamer((k,))
+    assert k0 != k1
+
+    square = gem.IndexSum(gem.Product(first(gather), second(gather)), (k0, k1))
+    result, = evaluate([square])
+    assert numpy.isclose(result.arr, (1.0 + 2.0 + 3.0) ** 2)

From e3790594d6072c63213602ef426baf62742083b2 Mon Sep 17 00:00:00 2001
From: Pablo Brubeck <brubeck@protonmail.com>
Date: Sat, 29 Aug 2026 10:37:39 +0100
Subject: [PATCH 10/11] Materialise a reduction view its consumer cannot index

inline_temporaries now offers a ComponentTensor over an IndexSum for inlining,
so that a consumer can index straight through the view rather than read a
temporary.  Only Indexed can do that: loopy lowers Indexed(ComponentTensor(e))
by reusing the temporary behind e.

Inverse, Solve and FlexiblyIndexed read the tensor whole, and a ComponentTensor
carries no name to give them.  Slate reaches all three -- A.solve(A*C) puts a
matrix-vector product under a Solve -- and dropping its Evaluate left codegen
with a node it cannot name:

    AttributeError: 'ComponentTensor' object has no attribute 'name'
    AssertionError: cannot generate expression from gem.gem.ComponentTensor

That accounted for 42 failures across tests/firedrake/slate.  Keep a reduction
view inlinable only where an Indexed consumes it.

Co-Authored-By: Claude Opus 5 <noreply@anthropic.com>
Claude-Session: https://claude.ai/code/session_013NoTXWyj2fVdJTHnMDFB4k
---
 gem/impero_utils.py | 9 +++++++++
 1 file changed, 9 insertions(+)

diff --git a/gem/impero_utils.py b/gem/impero_utils.py
index 2f15cd64..ea91ab17 100644
--- a/gem/impero_utils.py
+++ b/gem/impero_utils.py
@@ -133,6 +133,15 @@ def inline_temporaries(expressions, ops):
             if (expr.shape == () or reduction_view) and refcount[expr] == 1:
                 candidates.add(expr)
 
+    # A reduction view carries no name of its own, so it can only be inlined
+    # where the consumer indexes straight through it.  Inverse, Solve and
+    # FlexiblyIndexed read the tensor whole and need it materialised.
+    indexed = {child for node in traversal(expressions)
+               if isinstance(node, gem.Indexed)
+               for child in node.children}
+    candidates = {expr for expr in candidates
+                  if not isinstance(expr, gem.ComponentTensor) or expr in indexed}
+
     # Prevent inlining that pulls expressions into inner loops
     for node in traversal(expressions):
         for child in node.children:

From 3abb75b73e6ebc0d01c4872624372f1e5050c104 Mon Sep 17 00:00:00 2001
From: Pablo Brubeck <brubeck@protonmail.com>
Date: Sat, 29 Aug 2026 15:46:48 +0100
Subject: [PATCH 11/11] Pass a list where a sequence is expected

make_sum takes a sequence of summands.  Passing a Counter worked only
because iterating a Counter yields its keys.

Co-Authored-By: Claude Opus 5 <noreply@anthropic.com>
Claude-Session: https://claude.ai/code/session_013NoTXWyj2fVdJTHnMDFB4k
---
 gem/coffee.py | 2 +-
 1 file changed, 1 insertion(+), 1 deletion(-)

diff --git a/gem/coffee.py b/gem/coffee.py
index 418b4926..2cc4e1d9 100644
--- a/gem/coffee.py
+++ b/gem/coffee.py
@@ -236,7 +236,7 @@ def _extract_repeated_linear_maps(
                 if len(multiplicities) == 1:
                     multiplicity, = multiplicities
                     if multiplicity > 1:
-                        atomic = make_sum(counts)
+                        atomic = make_sum(list(counts))
                         factors.append(Literal(float(multiplicity)))
             atomics.append(atomic)
         result.add(