more typing work

Author: majosm

.basedpyright/baseline.json

@@ -14607,22 +14607,6 @@
                     "lineCount": 1
                 }
             },
-            {
-                "code": "reportPrivateUsage",
-                "range": {
-                    "startColumn": 53,
-                    "endColumn": 58,
-                    "lineCount": 1
-                }
-            },
-            {
-                "code": "reportPrivateUsage",
-                "range": {
-                    "startColumn": 53,
-                    "endColumn": 58,
-                    "lineCount": 1
-                }
-            },
             {
                 "code": "reportUnknownMemberType",
                 "range": {
@@ -14903,22 +14887,6 @@
                     "lineCount": 1
                 }
             },
-            {
-                "code": "reportPrivateUsage",
-                "range": {
-                    "startColumn": 53,
-                    "endColumn": 58,
-                    "lineCount": 1
-                }
-            },
-            {
-                "code": "reportPrivateUsage",
-                "range": {
-                    "startColumn": 53,
-                    "endColumn": 58,
-                    "lineCount": 1
-                }
-            },
             {
                 "code": "reportUnknownMemberType",
                 "range": {
@@ -20891,38 +20859,6 @@
                     "lineCount": 1
                 }
             },
-            {
-                "code": "reportAttributeAccessIssue",
-                "range": {
-                    "startColumn": 37,
-                    "endColumn": 54,
-                    "lineCount": 1
-                }
-            },
-            {
-                "code": "reportAttributeAccessIssue",
-                "range": {
-                    "startColumn": 25,
-                    "endColumn": 29,
-                    "lineCount": 1
-                }
-            },
-            {
-                "code": "reportAttributeAccessIssue",
-                "range": {
-                    "startColumn": 25,
-                    "endColumn": 29,
-                    "lineCount": 1
-                }
-            },
-            {
-                "code": "reportAttributeAccessIssue",
-                "range": {
-                    "startColumn": 38,
-                    "endColumn": 55,
-                    "lineCount": 1
-                }
-            },
             {
                 "code": "reportMissingParameterType",
                 "range": {

arraycontext/impl/pytato/__init__.py

@@ -152,9 +152,10 @@ def __init__(
         super().__init__()
 
         import pytato as pt
-        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,
+                pt.AbstractResultWithNamedArrays,
                 tuple[pt.AbstractResultWithNamedArrays, str]] = {}
 
         if compile_trace_callback is None:
@@ -600,12 +601,10 @@ def _to_frozen(key: tuple[Any, ...], ary) -> TaggableCLArray:
         dag = Deduplicator()(dag)
 
         # FIXME: Remove this if/when _normalize_pt_expr gets support for functions
-        dag = pt.tag_all_calls_to_be_inlined(
-            dag)
+        dag = pt.tag_all_calls_to_be_inlined(dag)
         dag = pt.inline_calls(dag)
 
-        normalized_expr, bound_arguments = _normalize_pt_expr(
-                dag)
+        normalized_expr, bound_arguments = _normalize_pt_expr(dag)
 
         try:
             pt_prg = self._freeze_prg_cache[normalized_expr]
@@ -760,10 +759,10 @@ def compile(self, f: Callable[..., Any]) -> Callable[..., Any]:
     def transform_dag(self, dag: pytato.AbstractResultWithNamedArrays
                       ) -> pytato.AbstractResultWithNamedArrays:
         import pytato as pt
-        tdag = pt.tag_all_calls_to_be_inlined(dag)
-        tdag = pt.inline_calls(tdag)
-        tdag = pt.transform.materialize_with_mpms(tdag)
-        return tdag
+        dag = pt.tag_all_calls_to_be_inlined(dag)
+        dag = pt.inline_calls(dag)
+        dag = pt.transform.materialize_with_mpms(dag)
+        return dag
 
     def einsum(self, spec, *args, arg_names=None, tagged=()):
         import pytato as pt

arraycontext/impl/pytato/utils.py

@@ -50,7 +50,6 @@
     Array,
     Axis as PtAxis,
     DataWrapper,
-    DictOfNamedArrays,
     Placeholder,
     SizeParam,
     make_placeholder,
@@ -138,7 +137,7 @@ def map_function_definition(
 # definitions can't contain non-argument placeholders
 def _normalize_pt_expr(
         expr: AbstractResultWithNamedArrays
-        ) -> tuple[DictOfNamedArrays, Mapping[str, Any]]:
+        ) -> tuple[AbstractResultWithNamedArrays, Mapping[str, Any]]:
     """
     Returns ``(normalized_expr, bound_arguments)``.  *normalized_expr* is a
     normalized form of *expr*, with all instances of
@@ -157,7 +156,6 @@ def _normalize_pt_expr(
 
     normalize_mapper = _DatawrapperToBoundPlaceholderMapper()
     normalized_expr = normalize_mapper(expr)
-    assert isinstance(normalized_expr, DictOfNamedArrays)
     return normalized_expr, normalize_mapper.bound_arguments