move _ary_container_key_stringifier to utils.py

Author: kaushikcfd

arraycontext/impl/pytato/__init__.py

@@ -436,9 +436,9 @@ def freeze(self, array):
         from arraycontext.container.traversal import rec_keyed_map_array_container
         from arraycontext.impl.pyopencl.taggable_cl_array import (
             TaggableCLArray, to_tagged_cl_array)
-        from arraycontext.impl.pytato.compile import _ary_container_key_stringifier
         from arraycontext.impl.pytato.utils import (
-            _normalize_pt_expr, get_cl_axes_from_pt_axes)
+            _ary_container_key_stringifier, _normalize_pt_expr,
+            get_cl_axes_from_pt_axes)
 
         array_as_dict: Dict[str, Union[cla.Array, TaggableCLArray, pt.Array]] = {}
         key_to_frozen_subary: Dict[str, TaggableCLArray] = {}
@@ -775,7 +775,7 @@ def freeze(self, array):
         import pytato as pt
 
         from arraycontext.container.traversal import rec_keyed_map_array_container
-        from arraycontext.impl.pytato.compile import _ary_container_key_stringifier
+        from arraycontext.impl.pytato.utils import _ary_container_key_stringifier
 
         array_as_dict: Dict[str, Union[jnp.ndarray, pt.Array]] = {}
         key_to_frozen_subary: Dict[str, jnp.ndarray] = {}

arraycontext/impl/pytato/compile.py

@@ -107,28 +107,6 @@ class LeafArrayDescriptor(AbstractInputDescriptor):
 
 # {{{ utilities
 
-def _ary_container_key_stringifier(keys: Tuple[Any, ...]) -> str:
-    """
-    Helper for :meth:`BaseLazilyCompilingFunctionCaller.__call__`. Stringifies an
-    array-container's component's key. Goals of this routine:
-
-    * No two different keys should have the same stringification
-    * Stringified key must a valid identifier according to :meth:`str.isidentifier`
-    * (informal) Shorter identifiers are preferred
-    """
-    def _rec_str(key: Any) -> str:
-        if isinstance(key, (str, int)):
-            return str(key)
-        elif isinstance(key, tuple):
-            # t in '_actx_t': stands for tuple
-            return "_actx_t" + "_".join(_rec_str(k) for k in key) + "_actx_endt"
-        else:
-            raise NotImplementedError("Key-stringication unimplemented for "
-                                      f"'{type(key).__name__}'.")
-
-    return "_".join(_rec_str(key) for key in keys)
-
-
 def _get_arg_id_to_arg_and_arg_id_to_descr(args: Tuple[Any, ...],
                                            kwargs: Mapping[str, Any]
                                            ) -> "Tuple[PMap[Tuple[Any, ...],\
@@ -318,6 +296,7 @@ def __call__(self, *args: Any, **kwargs: Any) -> Any:
         :attr:`~BaseLazilyCompilingFunctionCaller.f` with *args* in a lazy-sense.
         The intermediary pytato DAG for *args* is memoized in *self*.
         """
+        from arraycontext.impl.pytato.utils import _ary_container_key_stringifier
         arg_id_to_arg, arg_id_to_descr = _get_arg_id_to_arg_and_arg_id_to_descr(
             args, kwargs)
 

arraycontext/impl/pytato/utils.py

@@ -117,4 +117,30 @@ def get_loopy_target(self) -> Optional["lp.PyOpenCLTarget"]:
 
 # }}}
 
+
+# {{{ compile/outline helpers
+
+def _ary_container_key_stringifier(keys: Tuple[Any, ...]) -> str:
+    """
+    Helper for :meth:`BaseLazilyCompilingFunctionCaller.__call__`. Stringifies an
+    array-container's component's key. Goals of this routine:
+
+    * No two different keys should have the same stringification
+    * Stringified key must a valid identifier according to :meth:`str.isidentifier`
+    * (informal) Shorter identifiers are preferred
+    """
+    def _rec_str(key: Any) -> str:
+        if isinstance(key, (str, int)):
+            return str(key)
+        elif isinstance(key, tuple):
+            # t in '_actx_t': stands for tuple
+            return "_actx_t" + "_".join(_rec_str(k) for k in key) + "_actx_endt"
+        else:
+            raise NotImplementedError("Key-stringication unimplemented for "
+                                      f"'{type(key).__name__}'.")
+
+    return "_".join(_rec_str(key) for key in keys)
+
+# }}}
+
 # vim: foldmethod=marker

test/test_utils.py

@@ -34,7 +34,7 @@
 # {{{ test_pt_actx_key_stringification_uniqueness
 
 def test_pt_actx_key_stringification_uniqueness():
-    from arraycontext.impl.pytato.compile import _ary_container_key_stringifier
+    from arraycontext.impl.pytato.utils import _ary_container_key_stringifier
 
     assert (_ary_container_key_stringifier(((3, 2), 3))
             != _ary_container_key_stringifier((3, (2, 3))))