meshmode/discretization/connection/same_mesh.py

-from __future__ import division, print_function, absolute_import
-
 __copyright__ = "Copyright (C) 2014 Andreas Kloeckner"
 
 __license__ = """
@@ -24,26 +22,38 @@ THE SOFTWARE.
 
 import numpy as np
 
+from meshmode.transform_metadata import DiscretizationElementAxisTag
+
 
 # {{{ same-mesh constructor
 
 def make_same_mesh_connection(actx, to_discr, from_discr):
     from meshmode.discretization.connection.direct import (
-            InterpolationBatch,
-            DiscretizationConnectionElementGroup,
-            DirectDiscretizationConnection)
+        DirectDiscretizationConnection,
+        DiscretizationConnectionElementGroup,
+        IdentityDiscretizationConnection,
+        InterpolationBatch,
+    )
 
     if from_discr.mesh is not to_discr.mesh:
         raise ValueError("from_discr and to_discr must be based on "
                 "the same mesh")
 
+    if from_discr is to_discr:
+        return IdentityDiscretizationConnection(from_discr)
+
     groups = []
-    for igrp, (fgrp, tgrp) in enumerate(zip(from_discr.groups, to_discr.groups)):
-        all_elements = actx.freeze(
-                actx.from_numpy(
-                    np.arange(
-                        fgrp.nelements,
-                        dtype=np.intp)))
+    for igrp, (fgrp, tgrp) in enumerate(
+            zip(from_discr.groups, to_discr.groups, strict=True)):
+        from arraycontext.metadata import NameHint
+        all_elements = actx.tag(NameHint(f"all_el_ind_grp{igrp}"),
+                    actx.tag_axis(0,
+                        DiscretizationElementAxisTag(),
+                        actx.from_numpy(
+                            np.arange(
+                                fgrp.nelements,
+                                dtype=np.intp))))
+        all_elements = actx.freeze(all_elements)
         ibatch = InterpolationBatch(
                 from_group_index=igrp,
                 from_element_indices=all_elements,