add _verify_is_dag to fix pyright errors about assigning ArrayOrNames to DictOfNamedArrays

Author: majosm

arraycontext/impl/pytato/__init__.py

@@ -745,9 +745,13 @@ def compile(self, f: Callable[..., Any]) -> Callable[..., Any]:
     def transform_dag(self, dag: pytato.DictOfNamedArrays
                       ) -> pytato.DictOfNamedArrays:
         import pytato as pt
-        dag = pt.tag_all_calls_to_be_inlined(dag)
-        dag = pt.inline_calls(dag)
-        dag = pt.transform.materialize_with_mpms(dag)
+
+        # FIXME: Having to use _verify_is_dag seems clunky, but I'm not sure how to
+        # avoid it
+        from .utils import _verify_is_dag
+        dag = _verify_is_dag(pt.tag_all_calls_to_be_inlined(dag))
+        dag = _verify_is_dag(pt.inline_calls(dag))
+        dag = _verify_is_dag(pt.transform.materialize_with_mpms(dag))
         return dag
 
     def einsum(self, spec, *args, arg_names=None, tagged=()):

arraycontext/impl/pytato/outline.py

@@ -193,6 +193,8 @@ class OutlinedCall:
     def __call__(self, *args: Any, **kwargs: Any) -> ArrayOrContainer:
         arg_id_to_arg = _get_arg_id_to_arg(args, kwargs)
 
+        from .utils import _verify_is_dag
+
         if __debug__:
             # Add a prefix to the names to distinguish them from any existing
             # placeholders
@@ -201,9 +203,10 @@ def __call__(self, *args: Any, **kwargs: Any) -> ArrayOrContainer:
 
             prefixed_output = _call_with_placeholders(
                 self.f, args, kwargs, arg_id_to_prefixed_placeholder)
-            unpacked_prefixed_output = pt.transform.Deduplicator()(
-                pt.make_dict_of_named_arrays(
-                    _unpack_output(prefixed_output)))
+            unpacked_prefixed_output = _verify_is_dag(
+                pt.transform.Deduplicator()(
+                    pt.make_dict_of_named_arrays(
+                        _unpack_output(prefixed_output))))
 
             prefixed_placeholders = frozenset(
                 arg_id_to_prefixed_placeholder.values())
@@ -220,9 +223,10 @@ def __call__(self, *args: Any, **kwargs: Any) -> ArrayOrContainer:
         arg_id_to_placeholder = _get_arg_id_to_placeholder(arg_id_to_arg)
 
         output = _call_with_placeholders(self.f, args, kwargs, arg_id_to_placeholder)
-        unpacked_output = pt.transform.Deduplicator()(
-            pt.make_dict_of_named_arrays(
-                _unpack_output(output)))
+        unpacked_output = _verify_is_dag(
+            pt.transform.Deduplicator()(
+                pt.make_dict_of_named_arrays(
+                    _unpack_output(output))))
         if len(unpacked_output) == 1 and "_" in unpacked_output:
             ret_type = pt.function.ReturnType.ARRAY
         else:

arraycontext/impl/pytato/utils.py

@@ -72,6 +72,11 @@
     import loopy as lp
 
 
+def _verify_is_dag(dag: ArrayOrNames) -> DictOfNamedArrays:
+    assert isinstance(dag, DictOfNamedArrays)
+    return dag
+
+
 class _DatawrapperToBoundPlaceholderMapper(CopyMapper):
     """
     Helper mapper for :func:`normalize_pt_expr`. Every
@@ -141,7 +146,7 @@ def _normalize_pt_expr(
     Deterministic naming of placeholders permits more effective caching of
     equivalent graphs.
     """
-    expr = Deduplicator()(expr)
+    expr = _verify_is_dag(Deduplicator()(expr))
 
     if get_num_call_sites(expr):
         raise NotImplementedError(