Better transform typing

Author: inducer

arraycontext/impl/pytato/__init__.py

@@ -155,7 +155,7 @@ def __init__(
         self._freeze_prg_cache: dict[pt.DictOfNamedArrays, lp.TranslationUnit] = {}
         self._dag_transform_cache: dict[
                 pt.DictOfNamedArrays,
-                tuple[pt.DictOfNamedArrays, str]] = {}
+                tuple[pt.AbstractResultWithNamedArrays, str]] = {}
 
         if compile_trace_callback is None:
             def _compile_trace_callback(what, stage, ir):
@@ -177,8 +177,8 @@ def _frozen_array_types(self) -> tuple[type, ...]:
 
     # {{{ compilation
 
-    def transform_dag(self, dag: pytato.DictOfNamedArrays
-                      ) -> pytato.DictOfNamedArrays:
+    def transform_dag(self, dag: pytato.AbstractResultWithNamedArrays
+                      ) -> pytato.AbstractResultWithNamedArrays:
         """
         Returns a transformed version of *dag*. Sub-classes are supposed to
         override this method to implement context-specific transformations on
@@ -593,18 +593,19 @@ def _to_frozen(key: tuple[Any, ...], ary) -> TaggableCLArray:
                     rec_keyed_map_array_container(_to_frozen, array),
                     actx=None)
 
-        pt_dict_of_named_arrays = pt.make_dict_of_named_arrays(
+        dag = pt.make_dict_of_named_arrays(
                 key_to_pt_arrays)
 
-        pt_dict_of_named_arrays = pt.deduplicate(pt_dict_of_named_arrays)
+        from pytato.transform import Deduplicator
+        dag = Deduplicator()(dag)
 
         # FIXME: Remove this if/when _normalize_pt_expr gets support for functions
-        pt_dict_of_named_arrays = pt.tag_all_calls_to_be_inlined(
-            pt_dict_of_named_arrays)
-        pt_dict_of_named_arrays = pt.inline_calls(pt_dict_of_named_arrays)
+        dag = pt.tag_all_calls_to_be_inlined(
+            dag)
+        dag = pt.inline_calls(dag)
 
         normalized_expr, bound_arguments = _normalize_pt_expr(
-                pt_dict_of_named_arrays)
+                dag)
 
         try:
             pt_prg = self._freeze_prg_cache[normalized_expr]
@@ -756,13 +757,13 @@ def compile(self, f: Callable[..., Any]) -> Callable[..., Any]:
         from .compile import LazilyPyOpenCLCompilingFunctionCaller
         return LazilyPyOpenCLCompilingFunctionCaller(self, f)
 
-    def transform_dag(self, dag: pytato.DictOfNamedArrays
-                      ) -> pytato.DictOfNamedArrays:
+    def transform_dag(self, dag: pytato.AbstractResultWithNamedArrays
+                      ) -> pytato.AbstractResultWithNamedArrays:
         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)
-        return dag
+        tdag = pt.tag_all_calls_to_be_inlined(dag)
+        tdag = pt.inline_calls(tdag)
+        tdag = pt.transform.materialize_with_mpms(tdag)
+        return tdag
 
     def einsum(self, spec, *args, arg_names=None, tagged=()):
         import pytato as pt
@@ -977,8 +978,8 @@ def compile(self, f: Callable[..., Any]) -> Callable[..., Any]:
         return LazilyJAXCompilingFunctionCaller(self, f)
 
     @override
-    def transform_dag(self, dag: pytato.DictOfNamedArrays
-                      ) -> pytato.DictOfNamedArrays:
+    def transform_dag(self, dag: pytato.AbstractResultWithNamedArrays
+                      ) -> pytato.AbstractResultWithNamedArrays:
         import pytato as pt
 
         dag = pt.tag_all_calls_to_be_inlined(dag)

arraycontext/impl/pytato/compile.py

@@ -7,6 +7,9 @@
 """
 from __future__ import annotations
 
+from pytato.array import AxesT
+from pytato.transform import Deduplicator
+
 
 __copyright__ = """
 Copyright (C) 2020-1 University of Illinois Board of Trustees
@@ -189,7 +192,7 @@ def _to_input_for_compiled(
     """
     from arraycontext.impl.pyopencl.taggable_cl_array import to_tagged_cl_array
     if isinstance(ary, pt.Array):
-        dag = pt.deduplicate(pt.make_dict_of_named_arrays({"_actx_out": ary}))
+        dag = Deduplicator()(pt.make_dict_of_named_arrays({"_actx_out": ary}))
         # Transform the DAG to give metadata inference a chance to do its job
         return actx.transform_dag(dag)["_actx_out"].expr
     elif isinstance(ary, TaggableCLArray):
@@ -406,12 +409,16 @@ def _dag_to_transformed_pytato_prg(self, dict_of_named_arrays, *, prg_id=None):
         self.actx._compile_trace_callback(
                 prg_id, "post_transform_dag", pt_dict_of_named_arrays)
 
-        name_in_program_to_tags = {
-            name: out.tags
-            for name, out in pt_dict_of_named_arrays._data.items()}
-        name_in_program_to_axes = {
-            name: out.axes
-            for name, out in pt_dict_of_named_arrays._data.items()}
+        name_in_program_to_tags: dict[str, frozenset[Tag]] = {}
+        name_in_program_to_axes: dict[str, AxesT] = {}
+        if isinstance(pt_dict_of_named_arrays, pt.DictOfNamedArrays):
+            name_in_program_to_tags.update({
+                name: out.tags
+                for name, out in pt_dict_of_named_arrays._data.items()})
+
+            name_in_program_to_axes.update({
+                name: out.axes
+                for name, out in pt_dict_of_named_arrays._data.items()})
 
         self.actx._compile_trace_callback(
                 prg_id, "pre_generate_loopy", pt_dict_of_named_arrays)
@@ -503,12 +510,16 @@ def _dag_to_transformed_pytato_prg(self, dict_of_named_arrays, *, prg_id=None):
         self.actx._compile_trace_callback(
                 prg_id, "post_transform_dag", pt_dict_of_named_arrays)
 
-        name_in_program_to_tags = {
-            name: out.tags
-            for name, out in pt_dict_of_named_arrays._data.items()}
-        name_in_program_to_axes = {
-            name: out.axes
-            for name, out in pt_dict_of_named_arrays._data.items()}
+        name_in_program_to_tags: dict[str, frozenset[Tag]] = {}
+        name_in_program_to_axes: dict[str, AxesT] = {}
+        if isinstance(pt_dict_of_named_arrays, pt.DictOfNamedArrays):
+            name_in_program_to_tags.update({
+                name: out.tags
+                for name, out in pt_dict_of_named_arrays._data.items()})
+
+            name_in_program_to_axes.update({
+                name: out.axes
+                for name, out in pt_dict_of_named_arrays._data.items()})
 
         self.actx._compile_trace_callback(
                 prg_id, "pre_generate_jax", pt_dict_of_named_arrays)

arraycontext/impl/pytato/outline.py

@@ -1,5 +1,7 @@
 from __future__ import annotations
 
+from pytato.transform import Deduplicator
+
 
 __doc__ = """
 .. autoclass:: OutlinedCall
@@ -212,7 +214,7 @@ def __call__(self, *args: object, **kwargs: object) -> ArrayOrContainer:
 
             prefixed_output = _call_with_placeholders(
                 self.f, args, kwargs, arg_id_to_prefixed_placeholder)
-            unpacked_prefixed_output = pt.deduplicate(
+            unpacked_prefixed_output = Deduplicator()(
                 pt.make_dict_of_named_arrays(_unpack_output(prefixed_output)))
 
             prefixed_placeholders = frozenset(
@@ -230,7 +232,7 @@ def __call__(self, *args: object, **kwargs: object) -> 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.deduplicate(
+        unpacked_output = 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
@@ -247,10 +249,6 @@ def __call__(self, *args: object, **kwargs: object) -> ArrayOrContainer:
             for arg_id, placeholder in arg_id_to_placeholder.items()
             if placeholder in used_placeholders}
 
-        # pylint-disable-reason: pylint has a hard time with kw_only fields in
-        # dataclasses
-
-        # pylint: disable=unexpected-keyword-arg
         func_def = pt.function.FunctionDefinition(
             parameters=frozenset(call_bindings.keys()),
             return_type=ret_type,

arraycontext/impl/pytato/utils.py

@@ -44,6 +44,7 @@
 from typing_extensions import override
 
 import pytools
+from pytato import AbstractResultWithNamedArrays
 from pytato.analysis import get_num_call_sites
 from pytato.array import (
     Array,
@@ -59,8 +60,9 @@
 from pytato.transform import (
     ArrayOrNames,
     CopyMapper,
+    Deduplicator,
+    MappedT,
     TransformMapperCache,
-    deduplicate,
 )
 from pytools import UniqueNameGenerator, memoize_method
 
@@ -135,7 +137,7 @@ def map_function_definition(
 # FIXME: This strategy doesn't work if the DAG has functions, since function
 # definitions can't contain non-argument placeholders
 def _normalize_pt_expr(
-        expr: DictOfNamedArrays
+        expr: AbstractResultWithNamedArrays
         ) -> tuple[DictOfNamedArrays, Mapping[str, Any]]:
     """
     Returns ``(normalized_expr, bound_arguments)``.  *normalized_expr* is a
@@ -146,7 +148,7 @@ def _normalize_pt_expr(
     Deterministic naming of placeholders permits more effective caching of
     equivalent graphs.
     """
-    expr = deduplicate(expr)
+    expr = Deduplicator()(expr)
 
     if get_num_call_sites(expr):
         raise NotImplementedError(
@@ -246,15 +248,15 @@ def map_data_wrapper(self, expr: DataWrapper) -> Array:
             non_equality_tags=expr.non_equality_tags)
 
 
-def transfer_from_numpy(expr: ArrayOrNames, actx: ArrayContext) -> ArrayOrNames:
+def transfer_from_numpy(expr: MappedT, actx: ArrayContext) -> MappedT:
     """Transfer arrays contained in :class:`~pytato.array.DataWrapper`
     instances to be device arrays, using
     :meth:`~arraycontext.ArrayContext.from_numpy`.
     """
     return TransferFromNumpyMapper(actx)(expr)
 
 
-def transfer_to_numpy(expr: ArrayOrNames, actx: ArrayContext) -> ArrayOrNames:
+def transfer_to_numpy(expr: MappedT, actx: ArrayContext) -> MappedT:
     """Transfer arrays contained in :class:`~pytato.array.DataWrapper`
     instances to be :class:`numpy.ndarray` instances, using
     :meth:`~arraycontext.ArrayContext.to_numpy`.