tag a few more axes

inducer · Jun 22, 2023 · 8e0c041 · 8e0c041
1 parent e72ae4b
commit 8e0c041
Showing 1 changed file with 7 additions and 5 deletions.
diff --git a/meshmode/discretization/connection/direct.py b/meshmode/discretization/connection/direct.py
@@ -167,12 +167,14 @@ def _global_from_element_indices(
         np_full_from_element_indices[~np_from_el_present] = 0
 
         from_el_present = actx.freeze(
-                actx.tag(NameHint("from_el_present"),
-                    actx.from_numpy(
-                        np_from_el_present.astype(np.int8))))
+                actx.tag_axis(0, DiscretizationElementAxisTag(),
+                    actx.tag(NameHint("from_el_present"),
+                        actx.from_numpy(
+                            np_from_el_present.astype(np.int8)))))
         full_from_element_indices = actx.freeze(
-                actx.tag(NameHint("from_el_indices"),
-                    actx.from_numpy(np_full_from_element_indices)))
+                actx.tag_axis(0, DiscretizationElementAxisTag(),
+                    actx.tag(NameHint("from_el_indices"),
+                        actx.from_numpy(np_full_from_element_indices))))
 
         self._global_from_element_indices_cache = (
                 from_el_present, full_from_element_indices)