freeze constant arrays

inducer · Oct 13, 2022 · d27c898 · d27c898
1 parent 31d849e
commit d27c898
Showing 1 changed file with 7 additions and 2 deletions.
diff --git a/meshmode/discretization/connection/direct.py b/meshmode/discretization/connection/direct.py
@@ -232,6 +232,12 @@ class _FromGroupPickData(Generic[ArrayT]):
     from_element_indices: ArrayT
     is_surjective: bool
 
+    @keyed_memoize_method(key=lambda actx: type(actx).__name__)
+    def indexed_dof_pick_lists(self, actx):
+        assert actx.permits_advanced_indexing
+        return actx.freeze(
+            actx.thaw(self.dof_pick_lists)[actx.thaw(self.dof_pick_list_indices)])
+
 # }}}
 
 
@@ -737,8 +743,7 @@ def group_pick_knl(is_surjective: bool):
                             grp_ary_contrib = ary[fgpd.from_group_index][
                                         _reshape_and_preserve_tags(
                                             actx, from_element_indices, (-1, 1)),
-                                        actx.thaw(fgpd.dof_pick_lists)[
-                                            actx.thaw(fgpd.dof_pick_list_indices)]
+                                        actx.thaw(fgpd.indexed_dof_pick_lists(actx))
                                         ]
 
                             if not fgpd.is_surjective: