add ArrayOrNamesOrFunctionDef

Author: majosm

pytato/analysis/__init__.py

@@ -61,6 +61,7 @@
 from pytato.tags import ImplStored
 from pytato.transform import (
     ArrayOrNames,
+    ArrayOrNamesOrFunctionDef,
     ArrayOrNamesTc,
     CachedWalkMapper,
     CombineMapper,
@@ -362,7 +363,7 @@ def is_einsum_similar_to_subscript(expr: Einsum, subscripts: str) -> bool:
 
 class ListOfDirectPredecessorsGetter(
         Mapper[
-            list[ArrayOrNames | FunctionDefinition],
+            list[ArrayOrNamesOrFunctionDef],
             list[ArrayOrNames],
             []]):
     """
@@ -445,8 +446,8 @@ def map_distributed_send_ref_holder(self,
         return [expr.send.data, expr.passthrough_data]
 
     def map_call(
-            self, expr: Call) -> list[ArrayOrNames | FunctionDefinition]:
-        result: list[ArrayOrNames | FunctionDefinition] = []
+            self, expr: Call) -> list[ArrayOrNamesOrFunctionDef]:
+        result: list[ArrayOrNamesOrFunctionDef] = []
         if self.include_functions:
             result.append(expr.function)
         result += list(expr.bindings.values())
@@ -483,7 +484,7 @@ def __init__(self, *, include_functions: bool = False) -> None:
     @overload
     def __call__(
             self, expr: ArrayOrNames
-            ) -> FrozenOrderedSet[ArrayOrNames | FunctionDefinition]:
+            ) -> FrozenOrderedSet[ArrayOrNamesOrFunctionDef]:
         ...
 
     @overload
@@ -492,9 +493,9 @@ def __call__(self, expr: FunctionDefinition) -> FrozenOrderedSet[ArrayOrNames]:
 
     def __call__(
             self,
-            expr: ArrayOrNames | FunctionDefinition,
+            expr: ArrayOrNamesOrFunctionDef,
             ) -> (
-                FrozenOrderedSet[ArrayOrNames | FunctionDefinition]
+                FrozenOrderedSet[ArrayOrNamesOrFunctionDef]
                 | FrozenOrderedSet[ArrayOrNames]):
         """Get the direct predecessors of *expr*."""
         return FrozenOrderedSet(self._pred_getter(expr))
@@ -543,7 +544,7 @@ def clone_for_callee(self, function: FunctionDefinition) -> Self:
             _visited_functions=self._visited_functions)
 
     @override
-    def post_visit(self, expr: ArrayOrNames | FunctionDefinition) -> None:
+    def post_visit(self, expr: ArrayOrNamesOrFunctionDef) -> None:
         if not isinstance(expr, DictOfNamedArrays):
             self.expr_type_counts[type(expr)] += 1
 
@@ -606,7 +607,7 @@ def __init__(self, _visited_functions: set[Any] | None = None) -> None:
         super().__init__(_visited_functions=_visited_functions)
 
         self.expr_multiplicity_counts: \
-            dict[ArrayOrNames | FunctionDefinition, int] = defaultdict(int)
+            dict[ArrayOrNamesOrFunctionDef, int] = defaultdict(int)
 
     @override
     def get_cache_key(self, expr: ArrayOrNames) -> int:
@@ -619,13 +620,13 @@ def get_function_definition_cache_key(self, expr: FunctionDefinition) -> int:
         return id(expr)
 
     @override
-    def post_visit(self, expr: ArrayOrNames | FunctionDefinition) -> None:
+    def post_visit(self, expr: ArrayOrNamesOrFunctionDef) -> None:
         if not isinstance(expr, DictOfNamedArrays):
             self.expr_multiplicity_counts[expr] += 1
 
 
 def get_node_multiplicities(
-        outputs: ArrayOrNames) -> dict[ArrayOrNames | FunctionDefinition, int]:
+        outputs: ArrayOrNames) -> dict[ArrayOrNamesOrFunctionDef, int]:
     """
     Returns the multiplicity per `expr`.
     """
@@ -662,7 +663,7 @@ def get_function_definition_cache_key(self, expr: FunctionDefinition) -> int:
         return id(expr)
 
     @override
-    def post_visit(self, expr: ArrayOrNames | FunctionDefinition) -> None:
+    def post_visit(self, expr: ArrayOrNamesOrFunctionDef) -> None:
         if isinstance(expr, Call):
             self.count += 1
 
@@ -884,7 +885,7 @@ def map_call(self, expr: Call) -> None:
             f"{type(self).__name__} does not support functions.")
 
     @override
-    def post_visit(self, expr: ArrayOrNames | FunctionDefinition) -> None:
+    def post_visit(self, expr: ArrayOrNamesOrFunctionDef) -> None:
         if not is_materialized(expr):
             return
         assert isinstance(expr, Array)
@@ -978,7 +979,7 @@ def map_call(self, expr: Call) -> None:
             f"{type(self).__name__} does not support functions.")
 
     @override
-    def post_visit(self, expr: ArrayOrNames | FunctionDefinition) -> None:
+    def post_visit(self, expr: ArrayOrNamesOrFunctionDef) -> None:
         if not is_materialized(expr) or not has_taggable_materialization(expr):
             return
         assert isinstance(expr, Array)

pytato/distributed/verify.py

@@ -55,7 +55,11 @@
     DistributedGraphPartition,
     PartId,
 )
-from pytato.transform import ArrayOrNames, CachedWalkMapper
+from pytato.transform import (
+    ArrayOrNames,
+    ArrayOrNamesOrFunctionDef,
+    CachedWalkMapper,
+)
 
 
 logger = logging.getLogger(__name__)
@@ -68,7 +72,6 @@
     import numpy as np
 
     from pytato.distributed.nodes import CommTagType, DistributedRecv
-    from pytato.function import FunctionDefinition
 
 
 # {{{ data structures
@@ -156,7 +159,7 @@ def get_cache_key(self, expr: ArrayOrNames) -> int:
         return id(expr)
 
     @override
-    def visit(self, expr: ArrayOrNames | FunctionDefinition) -> bool:
+    def visit(self, expr: ArrayOrNamesOrFunctionDef) -> bool:
         super().visit(expr)
         if isinstance(expr, ArrayOrNames):
             self.seen_nodes.add(expr)

pytato/equality.py

@@ -63,6 +63,8 @@
 
 
 ArrayOrNames = Array | AbstractResultWithNamedArrays
+ArrayOrNamesOrFunctionDef = \
+    Array | AbstractResultWithNamedArrays | FunctionDefinition
 
 
 # {{{ EqualityComparer
@@ -87,7 +89,7 @@ def __init__(self) -> None:
         # Uses the same cache for both arrays and functions
         self._cache: dict[tuple[int, int], bool] = {}
 
-    def rec(self, expr1: ArrayOrNames | FunctionDefinition, expr2: object) -> bool:
+    def rec(self, expr1: ArrayOrNamesOrFunctionDef, expr2: object) -> bool:
         # These cases are simple enough that they don't need to be cached
         if expr1 is expr2:
             return True
@@ -119,7 +121,7 @@ def rec(self, expr1: ArrayOrNames | FunctionDefinition, expr2: object) -> bool:
             self._cache[cache_key] = result
             return result
 
-    def __call__(self, expr1: ArrayOrNames | FunctionDefinition, expr2: object) -> bool:
+    def __call__(self, expr1: ArrayOrNamesOrFunctionDef, expr2: object) -> bool:
         return self.rec(expr1, expr2)
 
     def handle_unsupported_array(self, expr1: Array,

pytato/transform/__init__.py

@@ -89,6 +89,8 @@
 
 
 ArrayOrNames: TypeAlias = Array | AbstractResultWithNamedArrays
+ArrayOrNamesOrFunctionDef: TypeAlias = \
+    Array | AbstractResultWithNamedArrays | FunctionDefinition
 ArrayOrNamesTc = TypeVar("ArrayOrNamesTc",
                   Array, AbstractResultWithNamedArrays, DictOfNamedArrays)
 ArrayOrNamesOrFunctionDefTc = TypeVar("ArrayOrNamesOrFunctionDefTc",
@@ -150,6 +152,7 @@
 ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
 
 .. class:: ArrayOrNames
+.. class:: ArrayOrNamesOrFunctionDef
 
 .. class:: ArrayOrNamesTc
 
@@ -307,7 +310,7 @@ def __call__(
 
     def __call__(
             self,
-            expr: ArrayOrNames | FunctionDefinition,
+            expr: ArrayOrNamesOrFunctionDef,
             *args: P.args,
             **kwargs: P.kwargs) -> ResultT | FunctionResultT:
         """Handle the mapping of *expr*."""
@@ -1569,7 +1572,7 @@ def clone_for_callee(
         return type(self)()
 
     def visit(
-            self, expr: ArrayOrNames | FunctionDefinition,
+            self, expr: ArrayOrNamesOrFunctionDef,
             *args: P.args, **kwargs: P.kwargs) -> bool:
         """
         If this method returns *True*, *expr* is traversed during the walk.
@@ -1579,7 +1582,7 @@ def visit(
         return True
 
     def post_visit(
-            self, expr: ArrayOrNames | FunctionDefinition,
+            self, expr: ArrayOrNamesOrFunctionDef,
             *args: P.args, **kwargs: P.kwargs) -> None:
         """
         Callback after *expr* has been traversed.
@@ -1841,7 +1844,7 @@ def __init__(
     def get_cache_key(self, expr: ArrayOrNames) -> int:
         return id(expr)
 
-    def post_visit(self, expr: ArrayOrNames | FunctionDefinition) -> None:
+    def post_visit(self, expr: ArrayOrNamesOrFunctionDef) -> None:
         if isinstance(expr, Array):
             self.topological_order.append(expr)
 

pytato/utils.py

@@ -61,16 +61,14 @@
     ScalarExpression,
     TypeCast,
 )
-from pytato.transform import ArrayOrNames, CachedMapper
+from pytato.transform import ArrayOrNamesOrFunctionDef, CachedMapper
 
 
 if TYPE_CHECKING:
     from collections.abc import Callable, Iterable, Sequence
 
     from pytools.tag import Tag
 
-    from pytato.function import FunctionDefinition
-
 
 __doc__ = """
 Helper routines
@@ -741,7 +739,7 @@ def get_einsum_specification(expr: Einsum) -> str:
     return f"{','.join(input_specs)}->{output_spec}"
 
 
-def is_materialized(expr: ArrayOrNames | FunctionDefinition) -> bool:
+def is_materialized(expr: ArrayOrNamesOrFunctionDef) -> bool:
     """Returns *True* if *expr* is materialized."""
     from pytato.array import InputArgumentBase
     from pytato.distributed.nodes import DistributedRecv
@@ -757,7 +755,7 @@ def is_materialized(expr: ArrayOrNames | FunctionDefinition) -> bool:
                 DistributedRecv)))
 
 
-def has_taggable_materialization(expr: ArrayOrNames | FunctionDefinition) -> bool:
+def has_taggable_materialization(expr: ArrayOrNamesOrFunctionDef) -> bool:
     """
     Returns *True* if *expr* uses the :class:`pytato.tags.ImplStored` tag to
     determine whether or not it is materialized.

pytato/visualization/dot.py

@@ -59,7 +59,12 @@
 )
 from pytato.function import Call, FunctionDefinition, NamedCallResult
 from pytato.tags import FunctionIdentifier
-from pytato.transform import ArrayOrNames, CachedMapper, InputGatherer
+from pytato.transform import (
+    ArrayOrNames,
+    ArrayOrNamesOrFunctionDef,
+    CachedMapper,
+    InputGatherer,
+)
 
 
 if TYPE_CHECKING:
@@ -160,7 +165,7 @@ def emit_subgraph(sg: _SubgraphTree) -> None:
 class _DotNodeInfo:
     title: str
     fields: dict[str, Any]
-    edges: dict[str, ArrayOrNames | FunctionDefinition]
+    edges: dict[str, ArrayOrNamesOrFunctionDef]
 
 
 def stringify_tags(tags: frozenset[Tag | None]) -> str:
@@ -193,7 +198,7 @@ def get_common_dot_info(self, expr: Array) -> _DotNodeInfo:
                   "non_equality_tags": expr.non_equality_tags,
                   }
 
-        edges: dict[str, ArrayOrNames | FunctionDefinition] = {}
+        edges: dict[str, ArrayOrNamesOrFunctionDef] = {}
         return _DotNodeInfo(title, fields, edges)
 
     # type-ignore-reason: incompatible with supertype
@@ -297,7 +302,7 @@ def map_einsum(self, expr: Einsum) -> None:
         self.node_to_dot[expr] = info
 
     def map_dict_of_named_arrays(self, expr: DictOfNamedArrays) -> None:
-        edges: dict[str, ArrayOrNames | FunctionDefinition] = {}
+        edges: dict[str, ArrayOrNamesOrFunctionDef] = {}
         for name, val in expr._data.items():
             edges[name] = val
             self.rec(val)
@@ -308,7 +313,7 @@ def map_dict_of_named_arrays(self, expr: DictOfNamedArrays) -> None:
                 edges=edges)
 
     def map_loopy_call(self, expr: LoopyCall) -> None:
-        edges: dict[str, ArrayOrNames | FunctionDefinition] = {}
+        edges: dict[str, ArrayOrNamesOrFunctionDef] = {}
         for name, arg in expr.bindings.items():
             if isinstance(arg, Array):
                 edges[name] = arg