Summarize changes to support prediction

README.rst

@@ -33,7 +33,3 @@ Pytato is written to pose no particular restrictions on the version of numpy
 used for execution. To use mypy-based type checking on Pytato itself or
 packages using Pytato, numpy 1.20 or newer is required, due to the
 typing-based changes to numpy in that release.
-
-Furthermore, pytato now uses type promotion rules based on those in
-`numpy <https://numpy.org/devdocs/numpy_2_0_migration_guide.html#changes-to-numpy-data-type-promotion>`__ that should result in the same
-data types as the currently installed version of numpy.

pytato/analysis/__init__.py

@@ -31,6 +31,7 @@
 
 from pymbolic.mapper.optimize import optimize_mapper
 from pytools import memoize_method
+from loopy.tools import LoopyKeyBuilder
 
 from pytato.array import (
     Array,
@@ -326,37 +327,37 @@ class DirectPredecessorsGetter(Mapper[frozenset[ArrayOrNames], []]):
 
         We only consider the predecessors of a nodes in a data-flow sense.
     """
-    def _get_preds_from_shape(self, shape: ShapeType) -> frozenset[ArrayOrNames]:
-        return frozenset({dim for dim in shape if isinstance(dim, Array)})
+    def _get_preds_from_shape(self, shape: ShapeType) -> dict[Array, None]:
+        return dict.fromkeys(dim for dim in shape if isinstance(dim, Array))
 
-    def map_index_lambda(self, expr: IndexLambda) -> frozenset[ArrayOrNames]:
-        return (frozenset(expr.bindings.values())
+    def map_index_lambda(self, expr: IndexLambda) -> dict[Array, None]:
+        return (dict.fromkeys(expr.bindings.values())
                 | self._get_preds_from_shape(expr.shape))
 
-    def map_stack(self, expr: Stack) -> frozenset[ArrayOrNames]:
-        return (frozenset(expr.arrays)
+    def map_stack(self, expr: Stack) -> dict[Array, None]:
+        return (dict.fromkeys(expr.arrays)
                 | self._get_preds_from_shape(expr.shape))
 
-    def map_concatenate(self, expr: Concatenate) -> frozenset[ArrayOrNames]:
-        return (frozenset(expr.arrays)
+    def map_concatenate(self, expr: Concatenate) -> dict[Array, None]:
+        return (dict.fromkeys(expr.arrays)
                 | self._get_preds_from_shape(expr.shape))
 
-    def map_einsum(self, expr: Einsum) -> frozenset[ArrayOrNames]:
-        return (frozenset(expr.args)
+    def map_einsum(self, expr: Einsum) -> dict[Array, None]:
+        return (dict.fromkeys(expr.args)
                 | self._get_preds_from_shape(expr.shape))
 
-    def map_loopy_call_result(self, expr: NamedArray) -> frozenset[ArrayOrNames]:
+    def map_loopy_call_result(self, expr: NamedArray) -> dict[Array, None]:
         from pytato.loopy import LoopyCall, LoopyCallResult
         assert isinstance(expr, LoopyCallResult)
         assert isinstance(expr._container, LoopyCall)
-        return (frozenset(ary
+        return (dict.fromkeys(ary
                           for ary in expr._container.bindings.values()
                           if isinstance(ary, Array))
                 | self._get_preds_from_shape(expr.shape))
 
-    def _map_index_base(self, expr: IndexBase) -> frozenset[ArrayOrNames]:
-        return (frozenset([expr.array])
-                | frozenset(idx for idx in expr.indices
+    def _map_index_base(self, expr: IndexBase) -> dict[Array, None]:
+        return (dict.fromkeys([expr.array])
+                | dict.fromkeys(idx for idx in expr.indices
                             if isinstance(idx, Array))
                 | self._get_preds_from_shape(expr.shape))
 
@@ -365,34 +366,34 @@ def _map_index_base(self, expr: IndexBase) -> frozenset[ArrayOrNames]:
     map_non_contiguous_advanced_index = _map_index_base
 
     def _map_index_remapping_base(self, expr: IndexRemappingBase
-                                  ) -> frozenset[ArrayOrNames]:
-        return frozenset([expr.array])
+                                  ) -> dict[ArrayOrNames, None]:
+        return dict.fromkeys([expr.array])
 
     map_roll = _map_index_remapping_base
     map_axis_permutation = _map_index_remapping_base
     map_reshape = _map_index_remapping_base
 
-    def _map_input_base(self, expr: InputArgumentBase) -> frozenset[ArrayOrNames]:
+    def _map_input_base(self, expr: InputArgumentBase) -> dict[Array, None]:
         return self._get_preds_from_shape(expr.shape)
 
     map_placeholder = _map_input_base
     map_data_wrapper = _map_input_base
     map_size_param = _map_input_base
 
-    def map_distributed_recv(self, expr: DistributedRecv) -> frozenset[ArrayOrNames]:
+    def map_distributed_recv(self, expr: DistributedRecv) -> dict[Array, None]:
         return self._get_preds_from_shape(expr.shape)
 
     def map_distributed_send_ref_holder(self,
                                         expr: DistributedSendRefHolder
-                                        ) -> frozenset[ArrayOrNames]:
-        return frozenset([expr.passthrough_data])
+                                        ) -> dict[ArrayOrNames, None]:
+        return dict.fromkeys([expr.passthrough_data])
 
-    def map_call(self, expr: Call) -> frozenset[ArrayOrNames]:
-        return frozenset(expr.bindings.values())
+    def map_call(self, expr: Call) -> dict[ArrayOrNames, None]:
+        return dict.fromkeys(expr.bindings.values())
 
     def map_named_call_result(
-            self, expr: NamedCallResult) -> frozenset[ArrayOrNames]:
-        return frozenset([expr._container])
+            self, expr: NamedCallResult) -> dict[ArrayOrNames, None]:
+        return dict.fromkeys([expr._container])
 
 
 # }}}
@@ -565,4 +566,47 @@ def get_num_call_sites(outputs: Array | DictOfNamedArrays) -> int:
 
 # }}}
 
+
+# {{{ PytatoKeyBuilder
+
+class PytatoKeyBuilder(LoopyKeyBuilder):  # type: ignore[misc]
+    """A custom :class:`pytools.persistent_dict.KeyBuilder` subclass
+    for objects within :mod:`pytato`.
+    """
+
+    def update_for_ndarray(self, key_hash: Any, key: Any) -> None:
+        self.rec(key_hash, key.data.tobytes())
+
+    def update_for_TaggableCLArray(self, key_hash: Any, key: Any) -> None:
+        self.rec(key_hash, key.get())
+
+    def update_for_Array(self, key_hash: Any, key: Any) -> None:
+        # CL Array
+        self.rec(key_hash, key.get())
+
+    # update_for_BitwiseAnd = LoopyKeyBuilder.update_for_pymbolic_expression  # noqa: N815, E501
+    # update_for_BitwiseNot = LoopyKeyBuilder.update_for_pymbolic_expression  # noqa: N815, E501
+    # update_for_BitwiseOr = LoopyKeyBuilder.update_for_pymbolic_expression  # noqa: N815, E501
+    # update_for_BitwiseXor = LoopyKeyBuilder.update_for_pymbolic_expression  # noqa: N815, E501
+    # update_for_Call = LoopyKeyBuilder.update_for_pymbolic_expression  # noqa: N815
+    # update_for_CallWithKwargs = LoopyKeyBuilder.update_for_pymbolic_expression  # noqa: N815, E501
+    # update_for_Comparison = LoopyKeyBuilder.update_for_pymbolic_expression  # noqa: N815, E501
+    # update_for_If = LoopyKeyBuilder.update_for_pymbolic_expression  # noqa: N815
+    # update_for_FloorDiv = LoopyKeyBuilder.update_for_pymbolic_expression  # noqa: N815, E501
+    # update_for_LeftShift = LoopyKeyBuilder.update_for_pymbolic_expression  # noqa: N815, E501
+    # update_for_LogicalAnd = LoopyKeyBuilder.update_for_pymbolic_expression  # noqa: N815, E501
+    # update_for_LogicalNot = LoopyKeyBuilder.update_for_pymbolic_expression  # noqa: N815, E501
+    # update_for_LogicalOr = LoopyKeyBuilder.update_for_pymbolic_expression  # noqa: N815, E501
+    # update_for_Lookup = LoopyKeyBuilder.update_for_pymbolic_expression  # noqa: N815, E501
+    # update_for_Power = LoopyKeyBuilder.update_for_pymbolic_expression  # noqa: N815
+    # update_for_Product = LoopyKeyBuilder.update_for_pymbolic_expression  # noqa: N815
+    # update_for_Quotient = LoopyKeyBuilder.update_for_pymbolic_expression  # noqa: N815, E501
+    # update_for_Remainder = LoopyKeyBuilder.update_for_pymbolic_expression  # noqa: N815, E501
+    # update_for_RightShift = LoopyKeyBuilder.update_for_pymbolic_expression  # noqa: N815, E501
+    # update_for_Subscript = LoopyKeyBuilder.update_for_pymbolic_expression  # noqa: N815, E501
+    # update_for_Sum = LoopyKeyBuilder.update_for_pymbolic_expression  # noqa: N815
+    # update_for_Variable = LoopyKeyBuilder.update_for_pymbolic_expression  # noqa: N815, E501
+
+# }}}
+
 # vim: fdm=marker

pytato/array.py

@@ -1,4 +1,5 @@
 from __future__ import annotations
+from traceback import FrameSummary, StackSummary
 
 
 __copyright__ = """
@@ -296,19 +297,25 @@ def normalize_shape_component(
 # {{{ array interface
 
 ConvertibleToIndexExpr = Union[int, slice, "Array", None, EllipsisType]
-IndexExpr = Union[IntegerT, "NormalizedSlice", "Array", None]
+IndexExpr = Union[IntegerT, "NormalizedSlice", "Array", None, EllipsisType]
+DtypeOrScalar = Union[_dtype_any, ScalarT]
+ArrayOrScalar = Union["Array", ScalarT]
 PyScalarType = type[bool] | type[int] | type[float] | type[complex]
 DtypeOrPyScalarType = _dtype_any | PyScalarType
 
 
-def _np_result_dtype(
-        *arrays_and_dtypes: np.typing.ArrayLike | np.typing.DTypeLike,
+# https://github.com/numpy/numpy/issues/19302
+def _np_result_type(
+        # actual dtype:
+        #*arrays_and_dtypes: Union[np.typing.ArrayLike, np.typing.DTypeLike],
+        # our dtype:
+        *arrays_and_dtypes: DtypeOrScalar,
         ) -> np.dtype[Any]:
     return np.result_type(*arrays_and_dtypes)
 
 
-def _truediv_result_type(*dtypes: DtypeOrPyScalarType) -> np.dtype[Any]:
-    dtype = _np_result_dtype(*dtypes)
+def _truediv_result_type(arg1: DtypeOrScalar, arg2: DtypeOrScalar) -> np.dtype[Any]:
+    dtype = _np_result_type(arg1, arg2)
     # See: test_true_divide in numpy/core/tests/test_ufunc.py
     # pylint: disable=no-member
     if dtype.kind in "iu":
@@ -571,11 +578,12 @@ def __matmul__(self, other: Array, reverse: bool = False) -> Array:
 
     def _binary_op(
                 self,
-                op: Callable[[ScalarExpression, ScalarExpression], ScalarExpression],
+                op: Callable[[ScalarExpression, ScalarExpression],
+                             ScalarExpression],
                 other: ArrayOrScalar,
                 get_result_type: Callable[
                         [ArrayOrScalar, ArrayOrScalar],
-                        np.dtype[Any]] = _np_result_dtype,
+                        np.dtype[Any]] = _np_result_type,
                 reverse: bool = False,
                 cast_to_result_dtype: bool = True,
                 is_pow: bool = False,
@@ -632,21 +640,33 @@ def _unary_op(self, op: Any) -> Array:
                 non_equality_tags=_get_created_at_tag(),
                 var_to_reduction_descr=immutabledict())
 
-    __mul__ = partialmethod(_binary_op, operator.mul)
-    __rmul__ = partialmethod(_binary_op, operator.mul, reverse=True)
+    # NOTE: Initializing the expression to "prim.Product(expr1, expr2)" is
+    # essential as opposed to performing "expr1 * expr2". This is to account
+    # for pymbolic's implementation of the "*" operator which might not
+    # instantiate the node corresponding to the operation when one of
+    # the operands is the neutral element of the operation.
+    #
+    # For the same reason 'prim.(Sum|FloorDiv|Quotient)' is preferred over the
+    # python operators on the operands.
 
-    __add__ = partialmethod(_binary_op, operator.add)
-    __radd__ = partialmethod(_binary_op, operator.add, reverse=True)
+    __mul__ = partialmethod(_binary_op, lambda l, r: prim.Product((l, r)))
+    __rmul__ = partialmethod(_binary_op, lambda l, r: prim.Product((l, r)),
+                             reverse=True)
 
-    __sub__ = partialmethod(_binary_op, operator.sub)
-    __rsub__ = partialmethod(_binary_op, operator.sub, reverse=True)
+    __add__ = partialmethod(_binary_op, lambda l, r: prim.Sum((l, r)))
+    __radd__ = partialmethod(_binary_op, lambda l, r: prim.Sum((l, r)),
+                             reverse=True)
 
-    __floordiv__ = partialmethod(_binary_op, operator.floordiv)
-    __rfloordiv__ = partialmethod(_binary_op, operator.floordiv, reverse=True)
+    __sub__ = partialmethod(_binary_op, lambda l, r: prim.Sum((l, -r)))
+    __rsub__ = partialmethod(_binary_op, lambda l, r: prim.Sum((l, -r)),
+                             reverse=True)
 
-    __truediv__ = partialmethod(_binary_op, operator.truediv,
+    __floordiv__ = partialmethod(_binary_op, prim.FloorDiv)
+    __rfloordiv__ = partialmethod(_binary_op, prim.FloorDiv, reverse=True)
+
+    __truediv__ = partialmethod(_binary_op, prim.Quotient,
             get_result_type=_truediv_result_type)
-    __rtruediv__ = partialmethod(_binary_op, operator.truediv,
+    __rtruediv__ = partialmethod(_binary_op, prim.Quotient,
             get_result_type=_truediv_result_type, reverse=True)
 
     __mod__ = partialmethod(_binary_op, operator.mod)
@@ -1397,7 +1417,7 @@ class Stack(_SuppliedAxesAndTagsMixin, Array):
 
     @property
     def dtype(self) -> np.dtype[Any]:
-        return _np_result_dtype(*(arr.dtype for arr in self.arrays))
+        return _np_result_type(*(arr.dtype for arr in self.arrays))
 
     @property
     def shape(self) -> ShapeType:
@@ -1430,7 +1450,7 @@ class Concatenate(_SuppliedAxesAndTagsMixin, Array):
 
     @property
     def dtype(self) -> np.dtype[Any]:
-        return _np_result_dtype(*(arr.dtype for arr in self.arrays))
+        return _np_result_type(*(arr.dtype for arr in self.arrays))
 
     @property
     def shape(self) -> ShapeType:
@@ -1547,6 +1567,7 @@ class Reshape(_SuppliedAxesAndTagsMixin, IndexRemappingBase):
 
     if __debug__:
         def __attrs_post_init__(self) -> None:
+            # assert self.non_equality_tags
             super().__attrs_post_init__()
 
     @property
@@ -2058,9 +2079,9 @@ def reshape(array: Array, newshape: int | Sequence[int],
         *and* the output array are linearized according to this order
         and 'matched up'.
 
-        Groups are found by multiplying axis lengths on the input and output side,
-        a matching input/output group is found once adding an input or axis to the
-        group makes the two products match.
+        Groups are found by multiplying axis lengths on the input and output
+        side, a matching input/output group is found once adding an input or
+        axis to the group makes the two products match.
 
         The semantics are identical to :func:`numpy.reshape`.