github/tinygrad
1
1
Fork 0
mirror of https://github.com/tinygrad/tinygrad.git synced 2026-04-29 03:00:14 -04:00
Code Issues Packages Projects Releases Wiki Activity

make SPEC=2 work for everything (#14476)

Browse Source 
* make SPEC=2 work for everything

* that's a horrible fix
This commit is contained in:
George Hotz
2026-02-02 10:30:56 +08:00
committed by GitHub
parent ea1f1d2b9d
commit 368a692e1a
2 changed files with 4 additions and 2 deletions
Download Patch File Download Diff File Expand all files Collapse all files

1
tinygrad/uop/ops.py
View File

@@ -1436,6 +1436,7 @@ def pyrender(ast:UOp) -> str:
for s in u.src: to_render.add(s)
if u.op is Ops.STORE: to_render.add(u.src[1])
if u.op in {Ops.REDUCE, Ops.REDUCE_AXIS}: to_render.add(u.src[0])
if u.op in {Ops.CUSTOM_KERNEL, Ops.CALL}: raise NotImplementedError("custom_kernel / call can't be pyrendered")
if u.op in not_rendered: continue
# checking the consumers is not enough, you have to make sure it's not used twice by the one consumer
if len(cmap[u]) == 1 and len([x for x in list(cmap[u].keys())[0].src if x is u]) == 1 and u.op not in always_rendered: continue

5
tinygrad/uop/spec.py
View File

@@ -313,8 +313,9 @@ def eval_pyrender(code:str) -> UOp:
return lcls['ast']
def test_pyrender(test_ast:UOp, assert_parents=True):
code = pyrender(test_ast)
ast:UOp = eval_pyrender(code)
try: code = pyrender(test_ast)
except NotImplementedError: return None # this is okay, not all ops can be pyrendered
ast = eval_pyrender(code)
if ast is not test_ast:
if assert_parents:
for u in test_ast.toposort(): test_pyrender(u, assert_parents=False)