tinygrad/tinygrad
2
0
Fork 3
mirror of https://github.com/tinygrad/tinygrad.git synced 2026-08-29 13:36:07 +00:00
Code Issues Packages Projects Releases Wiki Activity

fix resolve param (#15889)

Browse Source
This commit is contained in:
nimlgen
2026-04-23 17:41:44 +03:00
committed by GitHub
parent e4696185bd
commit 5cf4ad2fb6
2 changed files with 9 additions and 1 deletions
Download Patch File Download Diff File Expand all files Collapse all files
test/backend/test_multitensor.py
+8
View File
@@ -276,6 +276,14 @@ class TestMultiTensor(unittest.TestCase):
out = f(tt)
assert out.item() == 1+2+3+4
def test_multitensor_jit_input_reduce_shard_axis(self):
@TinyJit
def f(x): return x.sum(0).realize()
for _ in range(5):
tt = Tensor.ones(2, 64).contiguous().realize().shard((d1,d2), 0).realize()
out = f(tt)
np.testing.assert_allclose(out.numpy(), np.full(64, 2.0))
def test_multitensor_inside_jit(self):
@TinyJit
def f(x): return (x.shard((d1,d2), 0)+1).contiguous().sum()
tinygrad/engine/realize.py
+1 -1
View File
@@ -203,7 +203,7 @@ class ExecContext:
jit: bool = False
def _resolve(b:UOp, inputs:tuple[UOp, ...]) -> UOp:
if b.op is Ops.BUFFER_VIEW and b.src[0].op is Ops.PARAM: return b.replace(src=(inputs[b.src[0].arg], *b.src[1:]))
if b.op in (Ops.BUFFER_VIEW, Ops.MSELECT) and b.src[0].op is Ops.PARAM: return b.replace(src=(inputs[b.src[0].arg], *b.src[1:]))
return inputs[b.arg] if b.op is Ops.PARAM else b
def resolve_params(call:UOp, inputs:tuple[UOp, ...]) -> list[UOp]: return [_resolve(b, inputs) for b in call.src[1:] if b.op is not Ops.BIND]