mirrors/tinygrad
2
0
Fork 0
mirror of https://github.com/tinygrad/tinygrad.git synced 2026-06-24 02:14:17 +00:00
Code Issues Projects Releases Packages Wiki Activity

pickle schedules (#4114)

Browse source 
* pickle schedules

* Update test_pickle.py

* Update test_pickle.py

---------

Co-authored-by: George Hotz <72895+geohot@users.noreply.github.com>
This commit is contained in:
qazal 2024-04-09 23:47:25 +03:00 • committed by GitHub
commit 42edae8935
No known key found for this signature in database
GPG key ID: B5690EEEBB952194
1 changed files with 9 additions and 0 deletions
Download patch file Download diff file Expand all files Collapse all files

9
test/test_pickle.py
View file

@ -1,6 +1,7 @@
import unittest, pickle
import numpy as np
from tinygrad import Tensor, TinyJit
from tinygrad.engine.schedule import create_schedule
class TestPickle(unittest.TestCase):
def test_pickle_realized_tensor(self):
@ -37,5 +38,13 @@ class TestPickle(unittest.TestCase):
out = add_fxn(x, y)
np.testing.assert_equal(out.numpy(), 3)
def test_pickle_schedule(self):
a = Tensor([1,2])
out = a + 2
sched = create_schedule([out.lazydata])
pk = pickle.dumps(sched)
sched_pk = pickle.loads(pk)
assert sched_pk[-1].ast == sched[-1].ast
if __name__ == '__main__':
unittest.main()