fix beam in python 3.14+ (#13836)

* fix beam search on python 3.14

* add PickleableCount class to helpers

* change name, add test, add step

* tidy count init
This commit is contained in:
anu
2025-12-27 16:24:22 -05:00
committed by GitHub
parent 0f74909ae9
commit 9b4de8abc7
3 changed files with 24 additions and 4 deletions

View File

@@ -1,6 +1,6 @@
import ctypes, gzip, unittest, timeit
import ctypes, gzip, unittest, timeit, pickle
from tinygrad import Variable
from tinygrad.helpers import Context, ContextVar, argfix, colored, word_wrap, is_numpy_ndarray, mv_address, get_contraction
from tinygrad.helpers import Context, ContextVar, argfix, colored, word_wrap, is_numpy_ndarray, mv_address, get_contraction, count
from tinygrad.helpers import merge_dicts, strip_parens, prod, round_up, fetch, fully_flatten, from_mv, to_mv, polyN, time_to_str, cdiv, cmod, getbits
from tinygrad.tensor import Tensor, get_shape
import numpy as np
@@ -120,6 +120,18 @@ class TestRoundUp(unittest.TestCase):
self.assertEqual(round_up(232, 24984), 24984)
self.assertEqual(round_up(24984, 232), 25056)
class TestCount(unittest.TestCase):
def test_count_basic(self):
c = count(3)
self.assertEqual(next(c), 3)
self.assertEqual(next(c), 4)
def test_count_step_pickle(self):
c = count(1, 2)
self.assertEqual(next(c), 1)
c2 = pickle.loads(pickle.dumps(c))
self.assertEqual(next(c2), 3)
@unittest.skip("no fetch tests because they need internet")
class TestFetch(unittest.TestCase):
def test_fetch_bad_http(self):