From b28787583f5319c9a2f31babddd4b4ae85e11699 Mon Sep 17 00:00:00 2001
From: Matthias Diener <mdiener@illinois.edu>
Date: Tue, 20 Jun 2023 17:13:18 -0500
Subject: [PATCH] Numpy actx: better freeze/thaw
---
arraycontext/impl/numpy/__init__.py | 12 ++++++++++--
1 file changed, 10 insertions(+), 2 deletions(-)
diff --git a/arraycontext/impl/numpy/__init__.py b/arraycontext/impl/numpy/__init__.py
index dbc725f..89c4e88 100644
--- a/arraycontext/impl/numpy/__init__.py
+++ b/arraycontext/impl/numpy/__init__.py
@@ -38,6 +38,8 @@ import loopy as lp
from pytools.tag import Tag
from arraycontext.context import ArrayContext
+from arraycontext.container.traversal import (
+ rec_map_array_container, with_array_context)
class NumpyArrayContext(ArrayContext):
@@ -91,10 +93,16 @@ class NumpyArrayContext(ArrayContext):
return result
def freeze(self, array):
- return array
+ def _freeze(ary):
+ return ary
+
+ return with_array_context(rec_map_array_container(_freeze, array), actx=None)
def thaw(self, array):
- return array
+ def _thaw(ary):
+ return ary
+
+ return with_array_context(rec_map_array_container(_thaw, array), actx=self)
# }}}
--
GitLab