From 8d6cb4b69f8d976f53fb4dc0ff84dafbb139a980 Mon Sep 17 00:00:00 2001
From: Alexandru Fikl <alexfikl@gmail.com>
Date: Wed, 19 Aug 2026 20:26:29 +0300
Subject: [PATCH 1/2] feat: make flatten work on Operands

---
 pytential/symbolic/execution.py |  2 +-
 pytential/symbolic/mappers.py   | 13 +++++++++++--
 2 files changed, 12 insertions(+), 3 deletions(-)

diff --git a/pytential/symbolic/execution.py b/pytential/symbolic/execution.py
index e71389f7e..96ffb2941 100644
--- a/pytential/symbolic/execution.py
+++ b/pytential/symbolic/execution.py
@@ -638,7 +638,7 @@ def _prepare_expr(places: GeometryCollection,
 
     # FIXME: There's some mismatch between OperandTc and a conditional union
     # type that I was too impatient to figure out in detail.
-    expr = componentwise(flatten, expr)  # pyright: ignore[reportAssignmentType]
+    expr = flatten(expr)
     auto_source, auto_target = _prepare_auto_where(auto_where, places=places)
     expr = componentwise(  # pyright: ignore[reportAssignmentType]
             ToTargetTagger(
diff --git a/pytential/symbolic/mappers.py b/pytential/symbolic/mappers.py
index 5713740be..f2323b11f 100644
--- a/pytential/symbolic/mappers.py
+++ b/pytential/symbolic/mappers.py
@@ -380,8 +380,17 @@ class FlattenMapper(FlattenMapperBase, IdentityMapper):
     pass
 
 
-def flatten(expr: ArithmeticExpression):
-    return FlattenMapper().rec_arith(expr)
+def flatten(expr: pp.OperandTc) -> pp.OperandTc:
+    from pymbolic.geometric_algebra import MultiVector, componentwise
+    from pytools.obj_array import ObjectArray
+
+    def func(expr_i: ArithmeticExpression) -> ArithmeticExpression:
+        return FlattenMapper().rec_arith(expr_i)
+
+    if isinstance(expr, (ObjectArray, MultiVector)):
+        return componentwise(func, expr)
+    else:
+        return cast("pp.OperandTc", func(expr))
 
 # }}}
 

From f656e3bf34b9ed1515d3089ec7c47a6569397425 Mon Sep 17 00:00:00 2001
From: Alexandru Fikl <alexfikl@gmail.com>
Date: Wed, 19 Aug 2026 20:28:05 +0300
Subject: [PATCH 2/2] chore: update baseline

---
 .basedpyright/baseline.json | 16 ++++++++++++++++
 1 file changed, 16 insertions(+)

diff --git a/.basedpyright/baseline.json b/.basedpyright/baseline.json
index 58ef71ffa..6e69b5852 100644
--- a/.basedpyright/baseline.json
+++ b/.basedpyright/baseline.json
@@ -27913,6 +27913,14 @@
                     "lineCount": 1
                 }
             },
+            {
+                "code": "reportAbstractUsage",
+                "range": {
+                    "startColumn": 19,
+                    "endColumn": 18,
+                    "lineCount": 4
+                }
+            },
             {
                 "code": "reportCallIssue",
                 "range": {
@@ -28041,6 +28049,14 @@
                     "lineCount": 1
                 }
             },
+            {
+                "code": "reportAbstractUsage",
+                "range": {
+                    "startColumn": 19,
+                    "endColumn": 18,
+                    "lineCount": 4
+                }
+            },
             {
                 "code": "reportCallIssue",
                 "range": {