make the arange test check correctness [run_process_replay] (#5920)

This commit is contained in:
George Hotz 2024-08-05 13:41:06 -07:00 committed by GitHub
parent 8d1c884e78
commit e81c18f494
No known key found for this signature in database
GPG Key ID: B5690EEEBB952194
1 changed files with 5 additions and 1 deletions

View File

@ -4,11 +4,13 @@ from tinygrad import Tensor, GlobalCounters, dtypes
from tinygrad.helpers import Context, getenv
from tinygrad.engine.realize import run_schedule
from tinygrad.codegen.kernel import Opt, OptOps, Kernel
from tinygrad.engine.realize import CompiledRunner, ExecItem
class TestArange(unittest.TestCase):
def _get_flops(self, N, opts=None):
GlobalCounters.reset()
sched = Tensor.arange(N).schedule()
tt = Tensor.arange(N)
sched = tt.schedule()
self.assertEqual(len(sched), 1)
k = Kernel(sched[-1].ast)
if opts is not None:
@ -16,6 +18,8 @@ class TestArange(unittest.TestCase):
p = k.to_program()
print(p.name)
print(p.src)
ExecItem(CompiledRunner(p), [tt.lazydata.buffer]).run()
np.testing.assert_equal(tt.numpy(), np.arange(N))
return p.op_estimate
def test_complexity(self, opts=None):