Misc typing fixes

Co-authored-by: Matt Smith <mjsmith6@illinois.edu>

Author: inducer

arraycontext/impl/pytato/__init__.py

@@ -163,10 +163,11 @@ def __init__(
         """
         super().__init__()
 
-        self._freeze_prg_cache: dict[pt.DictOfNamedArrays, lp.TranslationUnit] = {}
+        self._freeze_prg_cache: dict[
+            pt.AbstractResultWithNamedArrays, lp.TranslationUnit] = {}
         self._dag_transform_cache: dict[
-                pt.DictOfNamedArrays,
-                tuple[pt.DictOfNamedArrays, str]] = {}
+                pt.AbstractResultWithNamedArrays,
+                tuple[pt.AbstractResultWithNamedArrays, str]] = {}
 
         if compile_trace_callback is None:
             def _compile_trace_callback(what, stage, ir):
@@ -226,8 +227,8 @@ def _tag_axis(ary: ArrayOrScalar) -> ArrayOrScalar:
 
     # {{{ 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
@@ -278,11 +279,12 @@ def get_target(self):
 
     # }}}
 
+    @override
     def outline(self,
                 f: Callable[..., Any],
                 *,
                 id: Hashable | None = None,
-                tags: frozenset[Tag] = frozenset()
+                tags: frozenset[Tag] = frozenset()  # pyright: ignore[reportCallInDefaultInitializer]
                 ) -> Callable[..., Any]:
         from pytato.tags import FunctionIdentifier
 
@@ -620,12 +622,10 @@ def _to_frozen(
             pt.make_dict_of_named_arrays(key_to_pt_arrays))
 
         # 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)
+        normalized_expr, bound_arguments = _normalize_pt_expr(dag)
 
         try:
             pt_prg = self._freeze_prg_cache[normalized_expr]
@@ -771,8 +771,8 @@ 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)
@@ -971,8 +971,9 @@ def compile(self, f: Callable[..., Any]) -> Callable[..., Any]:
         from .compile import LazilyJAXCompilingFunctionCaller
         return LazilyJAXCompilingFunctionCaller(self, f)
 
-    def transform_dag(self, dag: pytato.DictOfNamedArrays
-                      ) -> pytato.DictOfNamedArrays:
+    @override
+    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)

arraycontext/impl/pytato/compile.py

@@ -66,6 +66,7 @@
     from collections.abc import Callable, Hashable, Mapping
 
     import pyopencl.array as cla
+    from pytato.array import AxesT
 
 AllowedArray: TypeAlias = "pt.Array | TaggableCLArray | cla.Array"
 AllowedArrayTc = TypeVar("AllowedArrayTc", pt.Array, TaggableCLArray, "cla.Array")
@@ -408,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)
@@ -505,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/utils.py

@@ -10,6 +10,14 @@
 ^^^^^^^^^^^^^^^^^^^^^^^^^^^
 
 .. autofunction:: tabulate_profiling_data
+
+References
+^^^^^^^^^^
+
+.. autoclass:: ArrayOrNamesTc
+
+    A constrained type variable binding to either
+    :class:`pytato.Array` or :class:`pytato.AbstractResultWithNames`.
 """
 
 
@@ -40,22 +48,22 @@
 
 from typing import TYPE_CHECKING, Any, cast
 
+from typing_extensions import override
+
 import pytools
 from pytato.analysis import get_num_call_sites
 from pytato.array import (
-    AbstractResultWithNamedArrays,
     Array,
     Axis as PtAxis,
     DataWrapper,
-    DictOfNamedArrays,
     Placeholder,
     SizeParam,
     make_placeholder,
 )
-from pytato.function import FunctionDefinition
 from pytato.target.loopy import LoopyPyOpenCLTarget
 from pytato.transform import (
     ArrayOrNames,
+    ArrayOrNamesTc,
     CopyMapper,
     TransformMapperCache,
     deduplicate,
@@ -69,6 +77,8 @@
     from collections.abc import Mapping
 
     import loopy as lp
+    from pytato import AbstractResultWithNamedArrays
+    from pytato.function import FunctionDefinition
 
     from arraycontext import ArrayContext
     from arraycontext.container import SerializationKey
@@ -98,6 +108,7 @@ def __init__(
         self.vng = UniqueNameGenerator()
         self.seen_inputs: set[str] = set()
 
+    @override
     def map_data_wrapper(self, expr: DataWrapper) -> Array:
         if expr.name is not None:
             if expr.name in self.seen_inputs:
@@ -119,13 +130,16 @@ def map_data_wrapper(self, expr: DataWrapper) -> Array:
                     axes=expr.axes,
                     tags=expr.tags)
 
+    @override
     def map_size_param(self, expr: SizeParam) -> Array:
         raise NotImplementedError
 
+    @override
     def map_placeholder(self, expr: Placeholder) -> Array:
         raise ValueError("Placeholders cannot appear in"
                          " DatawrapperToBoundPlaceholderMapper.")
 
+    @override
     def map_function_definition(
             self, expr: FunctionDefinition) -> FunctionDefinition:
         raise ValueError("Function definitions cannot appear in"
@@ -135,8 +149,8 @@ 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
-        ) -> tuple[Array | AbstractResultWithNamedArrays, Mapping[str, Any]]:
+        expr: AbstractResultWithNamedArrays
+        ) -> tuple[AbstractResultWithNamedArrays, Mapping[str, Any]]:
     """
     Returns ``(normalized_expr, bound_arguments)``.  *normalized_expr* is a
     normalized form of *expr*, with all instances of
@@ -155,7 +169,6 @@ def _normalize_pt_expr(
 
     normalize_mapper = _DatawrapperToBoundPlaceholderMapper()
     normalized_expr = normalize_mapper(expr)
-    assert isinstance(normalized_expr, AbstractResultWithNamedArrays)
     return normalized_expr, normalize_mapper.bound_arguments
 
 
@@ -193,6 +206,7 @@ def __init__(self, actx: ArrayContext) -> None:
         super().__init__()
         self.actx = actx
 
+    @override
     def map_data_wrapper(self, expr: DataWrapper) -> Array:
         import numpy as np
 
@@ -225,6 +239,7 @@ def __init__(self, actx: ArrayContext) -> None:
         super().__init__()
         self.actx = actx
 
+    @override
     def map_data_wrapper(self, expr: DataWrapper) -> Array:
         import numpy as np
 
@@ -244,15 +259,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: ArrayOrNamesTc, actx: ArrayContext) -> ArrayOrNamesTc:
     """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: ArrayOrNamesTc, actx: ArrayContext) -> ArrayOrNamesTc:
     """Transfer arrays contained in :class:`~pytato.array.DataWrapper`
     instances to be :class:`numpy.ndarray` instances, using
     :meth:`~arraycontext.ArrayContext.to_numpy`.