mirror of https://github.com/commaai/tinygrad.git
35 lines
1.1 KiB
Python
35 lines
1.1 KiB
Python
import unittest
|
|
from tinygrad.tensor import Tensor
|
|
from tinygrad.engine.schedule import create_schedule
|
|
|
|
# stuff needed to unpack a kernel
|
|
# ruff: noqa: F401
|
|
from tinygrad.ops import MetaOps, LazyOp, TernaryOps, BinaryOps, UnaryOps, ReduceOps, BufferOps, MemBuffer, ConstBuffer
|
|
from tinygrad.lazy import LazyBuffer
|
|
from tinygrad import dtypes
|
|
from tinygrad.shape.shapetracker import ShapeTracker
|
|
from tinygrad.shape.view import View
|
|
from tinygrad.shape.symbolic import Variable
|
|
import numpy as np
|
|
import time
|
|
inf, nan = float('inf'), float('nan')
|
|
|
|
class TestLazyOp(unittest.TestCase):
|
|
def test_lazyop_str(self):
|
|
t = Tensor.rand(10) + Tensor.rand(10)
|
|
s = create_schedule([t.lazydata])
|
|
ast = s[-1].ast
|
|
ast_remade = eval(str(ast))
|
|
self.assertEqual(ast, ast_remade)
|
|
|
|
def test_selfreferential_speed(self):
|
|
st = time.monotonic()
|
|
for i in range(25):
|
|
p = Tensor([1]).lazydata
|
|
for _ in range(i): p = p.e(BinaryOps.ADD, p)
|
|
# sanity check if caching works this should be way faster
|
|
assert time.monotonic() -st < 0.5, f"{i}"
|
|
|
|
if __name__ == '__main__':
|
|
unittest.main()
|