From 3b9dd3330c58d638d6474aaa44391d63fd2df964 Mon Sep 17 00:00:00 2001
From: chenyu <chenyu@fastmail.com>
Date: Thu, 16 Nov 2023 18:35:08 -0500
Subject: [PATCH] add device to beam search cache key (#2333)

---
 tinygrad/features/search.py | 2 +-
 tinygrad/helpers.py         | 2 +-
 2 files changed, 2 insertions(+), 2 deletions(-)

diff --git a/tinygrad/features/search.py b/tinygrad/features/search.py
index bfe291c505..294392bc96 100644
--- a/tinygrad/features/search.py
+++ b/tinygrad/features/search.py
@@ -103,7 +103,7 @@ def get_linearizer_actions(lin:Linearizer, include_0=True) -> Dict[int, Lineariz
 def tuplize_uops(uops:List[UOp]) -> Tuple: return tuple([(x.uop, x.dtype, tuple(uops.index(x) for x in x.vin), x.arg) for x in uops])
 
 def beam_search(lin:Linearizer, rawbufs, amt:int, allow_test_size=True) -> Linearizer:
-  key = {"ast": str(lin.ast), "amt": amt, "allow_test_size": allow_test_size}
+  key = {"ast": str(lin.ast), "amt": amt, "allow_test_size": allow_test_size, "device": Device.DEFAULT}
   if (val:=diskcache_get("beam_search", key)) is not None and not getenv("IGNORE_BEAM_CACHE") and CACHELEVEL >= 1:
     ret = lin.copy()
     for o in val[len(lin.applied_opts):]: ret.apply_opt(o)
diff --git a/tinygrad/helpers.py b/tinygrad/helpers.py
index 86557d16d9..fa99c5edd1 100644
--- a/tinygrad/helpers.py
+++ b/tinygrad/helpers.py
@@ -173,7 +173,7 @@ _cache_dir: str = getenv("XDG_CACHE_HOME", os.path.expanduser("~/Library/Caches"
 CACHEDB: str = getenv("CACHEDB", os.path.abspath(os.path.join(_cache_dir, "tinygrad", "cache.db")))
 CACHELEVEL = getenv("CACHELEVEL", 2)
 
-VERSION = 6
+VERSION = 7
 _db_connection = None
 def db_connection():
   global _db_connection