avoid duplication in untag_loopy_call_results

Author: majosm

meshmode/array_context.py

@@ -262,9 +262,13 @@ def transform_dag(self, dag: pt_typ.DictOfNamedArrays) -> pt_typ.DictOfNamedArra
         def untag_loopy_call_results(
                     expr: pt.Array | pt.AbstractResultWithNamedArrays
                 ) -> pt.Array | pt.AbstractResultWithNamedArrays:
-            if isinstance(expr, pt.NamedArray):
-                return expr.copy(tags=frozenset(),
-                                 axes=(pt.Axis(frozenset()),)*expr.ndim)
+            if isinstance(expr, pt.loopy.LoopyCallResult):
+                new_tags = frozenset()
+                if any(axis.tags for axis in expr.axes):
+                    new_axes = (pt.Axis(frozenset()),)*expr.ndim
+                else:
+                    new_axes = expr.axes
+                return expr.replace_if_different(tags=new_tags, axes=new_axes)
             else:
                 return expr