CI: modernize the freezes on DOFArrays

inducer · May 9, 2022 · a430a74 · a430a74
1 parent f4bbf3e
commit a430a74
Showing 1 changed file with 4 additions and 7 deletions.
diff --git a/test/test_arraycontext.py b/test/test_arraycontext.py
@@ -32,7 +32,7 @@
         ArrayContext,
         dataclass_array_container, with_container_arithmetic,
         serialize_container, deserialize_container,
-        freeze, thaw,
+        freeze, thaw, with_array_context,
         FirstAxisIsElementsTag,
         PyOpenCLArrayContext,
         PytatoPyOpenCLArrayContext,
@@ -184,12 +184,9 @@ def _raise_index_inconsistency(i, stream_i):
                 for i, (stream_i, v) in enumerate(iterable)))
 
 
-@freeze.register(DOFArray)
-def _freeze_dofarray(ary, actx=None):
-    assert actx is None
-    return type(ary)(
-        None,
-        tuple(ary.array_context.freeze(subary) for subary in ary.data))
+@with_array_context.register(DOFArray)
+def _with_actx_dofarray(ary, actx):
+    return type(ary)(actx, ary.data)
 
 
 @thaw.register(DOFArray)