(hotifx) Seed types should be cacheable (#1623)

chaiNNer-org · Mar 3, 2023 · e6c7f0b · e6c7f0b
1 parent 13e0319
commit e6c7f0b
Show file tree

Hide file tree

Showing 2 changed files with 6 additions and 0 deletions.
diff --git a/backend/src/nodes/node_cache.py b/backend/src/nodes/node_cache.py
@@ -48,6 +48,9 @@ def _args_to_key(args) -> Tuple:
                 key.append(tuple(arg.shape))
                 key.append(arg.dtype.str)
                 key.append(hashlib.sha256(arg.tobytes()).digest())
+            elif hasattr(arg, "cache_key_func"):
+                key.append(arg.__class__.__name__)
+                key.append(arg.cache_key_func())
             else:
                 raise RuntimeError(f"Unexpected argument type {arg.__class__.__name__}")
         return tuple(key)

diff --git a/backend/src/nodes/utils/seed.py b/backend/src/nodes/utils/seed.py
@@ -30,3 +30,6 @@ def to_u32(self) -> int:
         Returns the value of the seed as a 32bit unsigned integer.
         """
         return self.to_range(0, _U32_MAX - 1)
+
+    def cache_key_func(self):
+        return self.value