Make flatten work on Operands

.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": {

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(

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))
 
 # }}}