Misc typing fixes

Author: inducer

arraycontext/impl/pytato/__init__.py

@@ -166,10 +166,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):
@@ -229,8 +230,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
@@ -635,12 +636,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]
@@ -786,8 +785,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)
@@ -986,8 +985,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`.
 """
 
 
@@ -38,24 +46,25 @@
 """
 
 
-from typing import TYPE_CHECKING, Any, cast
+from typing import TYPE_CHECKING, 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,
+    DataInterface,
     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 +78,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
@@ -94,10 +105,11 @@ def __init__(
             _cache=_cache,
             _function_cache=_function_cache)
 
-        self.bound_arguments: dict[str, Any] = {}
-        self.vng = UniqueNameGenerator()
+        self.bound_arguments: dict[str, DataInterface] = {}
+        self.vng: UniqueNameGenerator = 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 +131,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 +150,9 @@ 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, DataInterface]]:
     """
     Returns ``(normalized_expr, bound_arguments)``.  *normalized_expr* is a
     normalized form of *expr*, with all instances of
@@ -155,7 +171,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
 
 
@@ -172,7 +187,7 @@ def get_cl_axes_from_pt_axes(axes: tuple[PtAxis, ...]) -> tuple[ClAxis, ...]:
 class ArgSizeLimitingPytatoLoopyPyOpenCLTarget(LoopyPyOpenCLTarget):
     def __init__(self, limit_arg_size_nbytes: int) -> None:
         super().__init__()
-        self.limit_arg_size_nbytes = limit_arg_size_nbytes
+        self.limit_arg_size_nbytes: int = limit_arg_size_nbytes
 
     @memoize_method
     def get_loopy_target(self) -> lp.PyOpenCLTarget:
@@ -191,8 +206,9 @@ class TransferFromNumpyMapper(CopyMapper):
     """
     def __init__(self, actx: ArrayContext) -> None:
         super().__init__()
-        self.actx = actx
+        self.actx: ArrayContext = actx
 
+    @override
     def map_data_wrapper(self, expr: DataWrapper) -> Array:
         import numpy as np
 
@@ -223,8 +239,9 @@ class TransferToNumpyMapper(CopyMapper):
     """
     def __init__(self, actx: ArrayContext) -> None:
         super().__init__()
-        self.actx = actx
+        self.actx: ArrayContext = actx
 
+    @override
     def map_data_wrapper(self, expr: DataWrapper) -> Array:
         import numpy as np
 
@@ -244,15 +261,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`.
@@ -285,8 +302,7 @@ def tabulate_profiling_data(actx: PytatoPyOpenCLArrayContext) -> pytools.Table:
 
         t_sum = sum(times)
         t_avg = t_sum / num_calls
-        if t_sum is not None:
-            total_time += t_sum
+        total_time += t_sum
 
         tbl.add_row((kernel_name, num_calls, f"{t_sum:{g}}", f"{t_avg:{g}}"))