add _verify_is_array to avoid the need for rec_ary

the latter inflates recursion depth

Author: majosm

pytato/distributed/partition.py

@@ -90,7 +90,13 @@
     DistributedSendRefHolder,
 )
 from pytato.scalar_expr import SCALAR_CLASSES
-from pytato.transform import ArrayOrNames, CachedWalkMapper, CombineMapper, CopyMapper
+from pytato.transform import (
+    ArrayOrNames,
+    CachedWalkMapper,
+    CombineMapper,
+    CopyMapper,
+    _verify_is_array,
+)
 
 
 if TYPE_CHECKING:
@@ -396,7 +402,7 @@ def _make_distributed_partition(
 
         for name, val in name_to_part_output.items():
             assert name not in name_to_output
-            name_to_output[name] = comm_replacer.rec_ary(val)
+            name_to_output[name] = _verify_is_array(comm_replacer.rec(val))
 
         comm_ids = part_comm_ids[part_id]
 

pytato/transform/__init__.py

@@ -168,6 +168,11 @@ class UnsupportedArrayError(ValueError):
 P = ParamSpec("P")
 
 
+def _verify_is_array(expr: ArrayOrNames) -> Array:
+    assert isinstance(expr, Array)
+    return expr
+
+
 class Mapper(Generic[ResultT, FunctionResultT, P]):
     """A class that when called with a :class:`pytato.Array` recursively
     iterates over the DAG, calling the *_mapper_method* of each node. Users of
@@ -321,10 +326,7 @@ class TransformMapper(CachedMapper[ArrayOrNames, FunctionDefinition, []]):
     implement default mapper methods; for that, see :class:`CopyMapper`.
 
     """
-    def rec_ary(self, expr: Array) -> Array:
-        res = self.rec(expr)
-        assert isinstance(res, Array)
-        return res
+    pass
 
 # }}}
 
@@ -341,10 +343,7 @@ class TransformMapperWithExtraArgs(
     The logic in :class:`TransformMapper` purposely does not take the extra
     arguments to keep the cost of its each call frame low.
     """
-    def rec_ary(self, expr: Array, *args: P.args, **kwargs: P.kwargs) -> Array:
-        res = self.rec(expr, *args, **kwargs)
-        assert isinstance(res, Array)
-        return res
+    pass
 
 # }}}
 
@@ -390,33 +389,33 @@ def map_placeholder(self, expr: Placeholder) -> Array:
                 non_equality_tags=expr.non_equality_tags)
 
     def map_stack(self, expr: Stack) -> Array:
-        arrays = tuple(self.rec_ary(arr) for arr in expr.arrays)
+        arrays = tuple(_verify_is_array(self.rec(arr)) for arr in expr.arrays)
         return Stack(arrays=arrays, axis=expr.axis, axes=expr.axes, tags=expr.tags,
                 non_equality_tags=expr.non_equality_tags)
 
     def map_concatenate(self, expr: Concatenate) -> Array:
-        arrays = tuple(self.rec_ary(arr) for arr in expr.arrays)
+        arrays = tuple(_verify_is_array(self.rec(arr)) for arr in expr.arrays)
         return Concatenate(arrays=arrays, axis=expr.axis,
                            axes=expr.axes, tags=expr.tags,
                            non_equality_tags=expr.non_equality_tags)
 
     def map_roll(self, expr: Roll) -> Array:
-        return Roll(array=self.rec_ary(expr.array),
+        return Roll(array=_verify_is_array(self.rec(expr.array)),
                 shift=expr.shift,
                 axis=expr.axis,
                 axes=expr.axes,
                 tags=expr.tags,
                 non_equality_tags=expr.non_equality_tags)
 
     def map_axis_permutation(self, expr: AxisPermutation) -> Array:
-        return AxisPermutation(array=self.rec_ary(expr.array),
+        return AxisPermutation(array=_verify_is_array(self.rec(expr.array)),
                 axis_permutation=expr.axis_permutation,
                 axes=expr.axes,
                 tags=expr.tags,
                 non_equality_tags=expr.non_equality_tags)
 
     def _map_index_base(self, expr: IndexBase) -> Array:
-        return type(expr)(self.rec_ary(expr.array),
+        return type(expr)(_verify_is_array(self.rec(expr.array)),
                           indices=self.rec_idx_or_size_tuple(expr.indices),
                           axes=expr.axes,
                           tags=expr.tags,
@@ -453,7 +452,7 @@ def map_size_param(self, expr: SizeParam) -> Array:
 
     def map_einsum(self, expr: Einsum) -> Array:
         return Einsum(expr.access_descriptors,
-                      tuple(self.rec_ary(arg) for arg in expr.args),
+                      tuple(_verify_is_array(self.rec(arg)) for arg in expr.args),
                       axes=expr.axes,
                       redn_axis_to_redn_descr=expr.redn_axis_to_redn_descr,
                       tags=expr.tags,
@@ -470,7 +469,7 @@ def map_named_array(self, expr: NamedArray) -> Array:
 
     def map_dict_of_named_arrays(self,
             expr: DictOfNamedArrays) -> DictOfNamedArrays:
-        return DictOfNamedArrays({key: self.rec_ary(val.expr)
+        return DictOfNamedArrays({key: _verify_is_array(self.rec(val.expr))
                                   for key, val in expr.items()},
                                  tags=expr.tags
                                  )
@@ -498,7 +497,7 @@ def map_loopy_call_result(self, expr: LoopyCallResult) -> Array:
                 non_equality_tags=expr.non_equality_tags)
 
     def map_reshape(self, expr: Reshape) -> Array:
-        return Reshape(self.rec_ary(expr.array),
+        return Reshape(_verify_is_array(self.rec(expr.array)),
                        newshape=self.rec_idx_or_size_tuple(expr.newshape),
                        order=expr.order,
                        axes=expr.axes,
@@ -509,10 +508,10 @@ def map_distributed_send_ref_holder(
             self, expr: DistributedSendRefHolder) -> Array:
         return DistributedSendRefHolder(
                 send=DistributedSend(
-                    data=self.rec_ary(expr.send.data),
+                    data=_verify_is_array(self.rec(expr.send.data)),
                     dest_rank=expr.send.dest_rank,
                     comm_tag=expr.send.comm_tag),
-                passthrough_data=self.rec_ary(expr.passthrough_data),
+                passthrough_data=_verify_is_array(self.rec(expr.passthrough_data)),
                 )
 
     def map_distributed_recv(self, expr: DistributedRecv) -> Array:
@@ -590,19 +589,21 @@ def map_placeholder(self,
                            non_equality_tags=expr.non_equality_tags)
 
     def map_stack(self, expr: Stack, *args: P.args, **kwargs: P.kwargs) -> Array:
-        arrays = tuple(self.rec_ary(arr, *args, **kwargs) for arr in expr.arrays)
+        arrays = tuple(
+            _verify_is_array(self.rec(arr, *args, **kwargs)) for arr in expr.arrays)
         return Stack(arrays=arrays, axis=expr.axis, axes=expr.axes, tags=expr.tags,
                      non_equality_tags=expr.non_equality_tags)
 
     def map_concatenate(self,
                         expr: Concatenate, *args: P.args, **kwargs: P.kwargs) -> Array:
-        arrays = tuple(self.rec_ary(arr, *args, **kwargs) for arr in expr.arrays)
+        arrays = tuple(
+            _verify_is_array(self.rec(arr, *args, **kwargs)) for arr in expr.arrays)
         return Concatenate(arrays=arrays, axis=expr.axis,
                            axes=expr.axes, tags=expr.tags,
                            non_equality_tags=expr.non_equality_tags)
 
     def map_roll(self, expr: Roll, *args: P.args, **kwargs: P.kwargs) -> Array:
-        return Roll(array=self.rec_ary(expr.array, *args, **kwargs),
+        return Roll(array=_verify_is_array(self.rec(expr.array, *args, **kwargs)),
                     shift=expr.shift,
                     axis=expr.axis,
                     axes=expr.axes,
@@ -611,7 +612,8 @@ def map_roll(self, expr: Roll, *args: P.args, **kwargs: P.kwargs) -> Array:
 
     def map_axis_permutation(self, expr: AxisPermutation,
                              *args: P.args, **kwargs: P.kwargs) -> Array:
-        return AxisPermutation(array=self.rec_ary(expr.array, *args, **kwargs),
+        return AxisPermutation(array=_verify_is_array(
+                                   self.rec(expr.array, *args, **kwargs)),
                                axis_permutation=expr.axis_permutation,
                                axes=expr.axes,
                                tags=expr.tags,
@@ -620,7 +622,7 @@ def map_axis_permutation(self, expr: AxisPermutation,
     def _map_index_base(self,
                         expr: IndexBase, *args: P.args, **kwargs: P.kwargs) -> Array:
         assert isinstance(expr, _SuppliedAxesAndTagsMixin)
-        return type(expr)(self.rec_ary(expr.array, *args, **kwargs),
+        return type(expr)(_verify_is_array(self.rec(expr.array, *args, **kwargs)),
                           indices=self.rec_idx_or_size_tuple(expr.indices,
                                                              *args, **kwargs),
                           axes=expr.axes,
@@ -660,7 +662,8 @@ def map_size_param(self,
 
     def map_einsum(self, expr: Einsum, *args: P.args, **kwargs: P.kwargs) -> Array:
         return Einsum(expr.access_descriptors,
-                      tuple(self.rec_ary(arg, *args, **kwargs) for arg in expr.args),
+                      tuple(_verify_is_array(
+                          self.rec(arg, *args, **kwargs)) for arg in expr.args),
                       axes=expr.axes,
                       redn_axis_to_redn_descr=expr.redn_axis_to_redn_descr,
                       tags=expr.tags,
@@ -679,7 +682,8 @@ def map_named_array(self,
     def map_dict_of_named_arrays(self,
                 expr: DictOfNamedArrays, *args: P.args, **kwargs: P.kwargs
             ) -> DictOfNamedArrays:
-        return DictOfNamedArrays({key: self.rec_ary(val.expr, *args, **kwargs)
+        return DictOfNamedArrays({key: _verify_is_array(
+                                      self.rec(val.expr, *args, **kwargs))
                                   for key, val in expr.items()},
                                  tags=expr.tags,
                                  )
@@ -711,7 +715,7 @@ def map_loopy_call_result(self, expr: LoopyCallResult,
 
     def map_reshape(self, expr: Reshape,
                     *args: P.args, **kwargs: P.kwargs) -> Array:
-        return Reshape(self.rec_ary(expr.array, *args, **kwargs),
+        return Reshape(_verify_is_array(self.rec(expr.array, *args, **kwargs)),
                        newshape=self.rec_idx_or_size_tuple(expr.newshape,
                                                            *args, **kwargs),
                        order=expr.order,
@@ -723,10 +727,11 @@ def map_distributed_send_ref_holder(self, expr: DistributedSendRefHolder,
                                         *args: P.args, **kwargs: P.kwargs) -> Array:
         return DistributedSendRefHolder(
                 send=DistributedSend(
-                    data=self.rec_ary(expr.send.data, *args, **kwargs),
+                    data=_verify_is_array(self.rec(expr.send.data, *args, **kwargs)),
                     dest_rank=expr.send.dest_rank,
                     comm_tag=expr.send.comm_tag),
-                passthrough_data=self.rec_ary(expr.passthrough_data, *args, **kwargs))
+                passthrough_data=_verify_is_array(
+                    self.rec(expr.passthrough_data, *args, **kwargs)))
 
     def map_distributed_recv(self, expr: DistributedRecv,
                              *args: P.args, **kwargs: P.kwargs) -> Array:
@@ -1619,7 +1624,7 @@ def copy_dict_of_named_arrays(source_dict: DictOfNamedArrays,
     if not source_dict:
         data = {}
     else:
-        data = {name: copy_mapper.rec_ary(val.expr)
+        data = {name: _verify_is_array(copy_mapper.rec(val.expr))
                 for name, val in sorted(source_dict.items())}
 
     return DictOfNamedArrays(data, tags=source_dict.tags)

pytato/transform/calls.py

@@ -40,7 +40,7 @@
 )
 from pytato.function import Call, NamedCallResult
 from pytato.tags import InlineCallTag
-from pytato.transform import ArrayOrNames, CopyMapper
+from pytato.transform import ArrayOrNames, CopyMapper, _verify_is_array
 
 
 if TYPE_CHECKING:
@@ -79,7 +79,7 @@ def map_call(self, expr: Call) -> AbstractResultWithNamedArrays:
             substitutor = PlaceholderSubstitutor(new_expr.bindings)
 
             return DictOfNamedArrays(
-                {name: substitutor.rec_ary(ret)
+                {name: _verify_is_array(substitutor.rec(ret))
                  for name, ret in new_expr.function.returns.items()},
                 tags=new_expr.tags
             )

pytato/transform/einsum_distributive_law.py

@@ -59,6 +59,7 @@
 from pytato.transform import (
     MappedT,
     TransformMapperWithExtraArgs,
+    _verify_is_array,
 )
 from pytato.utils import are_shapes_equal
 
@@ -186,38 +187,42 @@ def map_index_lambda(self,
                         and isinstance(hlo.x2, Array)
                         and are_shapes_equal(hlo.x1.shape, hlo.x2.shape))
                 # https://github.com/python/mypy/issues/16499
-                return self.rec_ary(hlo.x1, ctx) + self.rec_ary(hlo.x2, ctx)  # type: ignore[no-any-return]
+                return (  # type: ignore[no-any-return]
+                    _verify_is_array(self.rec(hlo.x1, ctx))
+                    + _verify_is_array(self.rec(hlo.x2, ctx)))
             elif hlo.binary_op == BinaryOpType.SUB:
                 assert (isinstance(hlo.x1, Array)
                         and isinstance(hlo.x2, Array)
                         and are_shapes_equal(hlo.x1.shape, hlo.x2.shape))
                 assert are_shapes_equal(hlo.x1.shape, hlo.x2.shape)
                 # https://github.com/python/mypy/issues/16499
-                return self.rec_ary(hlo.x1, ctx) - self.rec_ary(hlo.x2, ctx)  # type: ignore[no-any-return]
+                return (  # type: ignore[no-any-return]
+                    _verify_is_array(self.rec(hlo.x1, ctx))
+                    - _verify_is_array(self.rec(hlo.x2, ctx)))
             elif hlo.binary_op == BinaryOpType.MULT:
                 if isinstance(hlo.x1, Array) and np.isscalar(hlo.x2):
                     # https://github.com/python/mypy/issues/16499
-                    return self.rec_ary(hlo.x1, ctx) * hlo.x2  # type: ignore[no-any-return]
+                    return _verify_is_array(self.rec(hlo.x1, ctx)) * hlo.x2  # type: ignore[no-any-return]
                 else:
                     assert isinstance(hlo.x2, Array) and np.isscalar(hlo.x1)
                     # https://github.com/python/mypy/issues/16499
-                    return hlo.x1 * self.rec_ary(hlo.x2, ctx)  # type: ignore[no-any-return]
+                    return hlo.x1 * _verify_is_array(self.rec(hlo.x2, ctx))  # type: ignore[no-any-return]
             elif hlo.binary_op == BinaryOpType.TRUEDIV:
                 if isinstance(hlo.x1, Array) and np.isscalar(hlo.x2):
                     # https://github.com/python/mypy/issues/16499
-                    return self.rec_ary(hlo.x1, ctx) / hlo.x2  # type: ignore[no-any-return]
+                    return _verify_is_array(self.rec(hlo.x1, ctx)) / hlo.x2  # type: ignore[no-any-return]
                 else:
                     assert isinstance(hlo.x2, Array) and np.isscalar(hlo.x1)
                     # https://github.com/python/mypy/issues/16499
-                    return hlo.x1 / self.rec_ary(hlo.x2, ctx)  # type: ignore[no-any-return]
+                    return hlo.x1 / _verify_is_array(self.rec(hlo.x2, ctx))  # type: ignore[no-any-return]
             else:
                 raise NotImplementedError(hlo)
         else:
             rec_expr = IndexLambda(
                 expr=expr.expr,
                 shape=expr.shape,
                 dtype=expr.dtype,
-                bindings=immutabledict({name: self.rec_ary(bnd, None)
+                bindings=immutabledict({name: _verify_is_array(self.rec(bnd, None))
                               for name, bnd in sorted(expr.bindings.items())}),
                 var_to_reduction_descr=expr.var_to_reduction_descr,
                 tags=expr.tags,
@@ -244,12 +249,13 @@ def map_einsum(self,
                     tags=expr.tags,
                     axes=expr.axes,
                 )
-                return self.rec_ary(expr.args[distributive_law_descr.ioperand], ctx)
+                return _verify_is_array(
+                    self.rec(expr.args[distributive_law_descr.ioperand], ctx))
         else:
             assert isinstance(distributive_law_descr, DoNotDistribute)
             rec_expr = Einsum(
                 expr.access_descriptors,
-                tuple(self.rec_ary(arg, None) for arg in expr.args),
+                tuple(_verify_is_array(self.rec(arg, None)) for arg in expr.args),
                 expr.redn_axis_to_redn_descr,
                 tags=expr.tags,
                 axes=expr.axes
@@ -260,7 +266,7 @@ def map_einsum(self,
     def map_stack(self,
                   expr: Stack,
                   ctx: _EinsumDistributiveLawMapperContext | None) -> Array:
-        rec_expr = Stack(tuple(self.rec_ary(ary, None)
+        rec_expr = Stack(tuple(_verify_is_array(self.rec(ary, None))
                                for ary in expr.arrays),
                          expr.axis,
                          tags=expr.tags,
@@ -271,7 +277,7 @@ def map_concatenate(self,
                         expr: Concatenate,
                         ctx: _EinsumDistributiveLawMapperContext | None
                         ) -> Array:
-        rec_expr = Concatenate(tuple(self.rec_ary(ary, None)
+        rec_expr = Concatenate(tuple(_verify_is_array(self.rec(ary, None))
                                      for ary in expr.arrays),
                                expr.axis,
                                tags=expr.tags,
@@ -282,7 +288,7 @@ def map_roll(self,
                  expr: Roll,
                  ctx: _EinsumDistributiveLawMapperContext | None
                  ) -> Array:
-        rec_expr = Roll(self.rec_ary(expr.array, None),
+        rec_expr = Roll(_verify_is_array(self.rec(expr.array, None)),
                         expr.shift,
                         expr.axis,
                         tags=expr.tags,
@@ -293,7 +299,7 @@ def map_axis_permutation(self,
                              expr: AxisPermutation,
                              ctx: _EinsumDistributiveLawMapperContext | None
                              ) -> Array:
-        rec_expr = AxisPermutation(self.rec_ary(expr.array, None),
+        rec_expr = AxisPermutation(_verify_is_array(self.rec(expr.array, None)),
                                    expr.axis_permutation,
                                    tags=expr.tags,
                                    axes=expr.axes)
@@ -303,7 +309,7 @@ def _map_index_base(self,
                         expr: IndexBase,
                         ctx: _EinsumDistributiveLawMapperContext | None
                         ) -> Array:
-        rec_expr = type(expr)(self.rec_ary(expr.array, None),
+        rec_expr = type(expr)(_verify_is_array(self.rec(expr.array, None)),
                               expr.indices,
                               tags=expr.tags,
                               axes=expr.axes)
@@ -317,7 +323,7 @@ def map_reshape(self,
                     expr: Reshape,
                     ctx: _EinsumDistributiveLawMapperContext | None
                     ) -> Array:
-        rec_expr = Reshape(self.rec_ary(expr.array, None),
+        rec_expr = Reshape(_verify_is_array(self.rec(expr.array, None)),
                            expr.newshape,
                            expr.order,
                            tags=expr.tags,

pytato/transform/remove_broadcasts_einsum.py

@@ -31,7 +31,7 @@
 from typing import cast
 
 from pytato.array import Array, Einsum, EinsumAxisDescriptor
-from pytato.transform import CopyMapper, MappedT
+from pytato.transform import CopyMapper, MappedT, _verify_is_array
 from pytato.utils import are_shape_components_equal
 
 
@@ -42,7 +42,7 @@ def map_einsum(self, expr: Einsum) -> Array:
         descr_to_axis_len = expr._access_descr_to_axis_len()
 
         for acc_descrs, arg in zip(expr.access_descriptors, expr.args, strict=True):
-            arg = self.rec_ary(arg)
+            arg = _verify_is_array(self.rec(arg))
             axes_to_squeeze: list[int] = []
             for idim, acc_descr in enumerate(acc_descrs):
                 if not are_shape_components_equal(arg.shape[idim],