add PytatoKeyBuilder

pytato/analysis/__init__.py

@@ -36,6 +36,7 @@
 from pytato.loopy import LoopyCall
 from pymbolic.mapper.optimize import optimize_mapper
 from pytools import memoize_method
+from loopy.tools import LoopyKeyBuilder, PersistentHashWalkMapper
 
 if TYPE_CHECKING:
     from pytato.distributed.nodes import DistributedRecv, DistributedSendRefHolder
@@ -453,4 +454,32 @@ def get_num_call_sites(outputs: Union[Array, DictOfNamedArrays]) -> int:
 
 # }}}
 
+
+# {{{ PytatoKeyBuilder
+
+class PytatoKeyBuilder(LoopyKeyBuilder):
+    """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, hash(key.data.tobytes()))  # type: ignore[no-untyped-call]
+
+    def update_for_pymbolic_expression(self, key_hash: Any, key: Any) -> None:
+        if key is None:
+            self.update_for_NoneType(key_hash, key)  # type: ignore[no-untyped-call]
+        else:
+            PersistentHashWalkMapper(key_hash)(key)
+
+    update_for_Product = update_for_pymbolic_expression  # noqa: N815
+    update_for_Sum = update_for_pymbolic_expression  # noqa: N815
+    update_for_If = update_for_pymbolic_expression  # noqa: N815
+    update_for_LogicalOr = update_for_pymbolic_expression  # noqa: N815
+    update_for_Call = update_for_pymbolic_expression  # noqa: N815
+    update_for_Comparison = update_for_pymbolic_expression  # noqa: N815
+    update_for_Quotient = update_for_pymbolic_expression  # noqa: N815
+    update_for_Power = update_for_pymbolic_expression  # noqa: N815
+
+# }}}
+
 # vim: fdm=marker

test/test_pytato.py

@@ -1116,6 +1116,26 @@ def test_dot_visualizers():
     # }}}
 
 
+def test_persistent_dict():
+    from pytools.persistent_dict import WriteOncePersistentDict, ReadOnlyEntryError
+    from pytato.analysis import PytatoKeyBuilder
+
+    axis_len = 5
+
+    pd = WriteOncePersistentDict("test_persistent_dict", key_builder=PytatoKeyBuilder(), container_dir="./pytest-pdict")
+
+    for i in range(100):
+        rdagc = RandomDAGContext(np.random.default_rng(seed=i),
+                axis_len=axis_len, use_numpy=True)
+
+        dag = make_random_dag(rdagc)
+        pd[dag] = 42
+
+        # Make sure key stays the same
+        with pytest.raises(ReadOnlyEntryError):
+            pd[dag] = 42
+
+
 if __name__ == "__main__":
     if len(sys.argv) > 1:
         exec(sys.argv[1])