mirror of
https://github.com/tinygrad/tinygrad.git
synced 2026-04-07 03:00:26 -04:00
a949de873b
* reduce 2.0 * works * hacks * DEBUG=3 for shapes * fix types * 0s weren't being folded * cleaner * last_reduce is no longer needed * comments and cleanup
25 lines
766 B
Python
25 lines
766 B
Python
import sys
|
|
import numpy as np
|
|
from typing import Dict
|
|
from tinygrad.ast import ASTKernel
|
|
from tinygrad.llops.ops_cpu import CPUBuffer
|
|
from tinygrad.ops import DeviceBuffer
|
|
from tinygrad.lazy import realize_buffers
|
|
|
|
in_test = False
|
|
test_cnt = 0
|
|
def test_ast(k:ASTKernel):
|
|
global in_test, test_cnt
|
|
if in_test: return
|
|
in_test = True
|
|
print("testing AST", test_cnt)
|
|
test_cnt += 1
|
|
cpubufs : Dict[DeviceBuffer, CPUBuffer] = {x:CPUBuffer.fromCPU(x.toCPU()) for x in k.bufs}
|
|
real_out = cpubufs[k.bufs[0]]
|
|
test_out = CPUBuffer.exec_ast(realize_buffers(cpubufs, k.ast))
|
|
if not np.allclose(real_out, test_out, atol=1e-4):
|
|
print("MISMATCH")
|
|
print(k.print())
|
|
sys.tracebacklimit = 0
|
|
np.testing.assert_allclose(real_out, test_out)
|
|
in_test = False |