From 09574a096abcafaf820b2f33c8850dd92bd35b37 Mon Sep 17 00:00:00 2001 From: George Hotz Date: Mon, 16 Mar 2026 12:04:39 +0800 Subject: [PATCH] maketuple --- test/unit/test_function.py | 28 ++++++++++++++++------------ tinygrad/engine/allocations.py | 2 +- tinygrad/function.py | 2 +- tinygrad/gradient.py | 2 +- tinygrad/uop/ops.py | 2 +- 5 files changed, 20 insertions(+), 16 deletions(-) diff --git a/test/unit/test_function.py b/test/unit/test_function.py index 32c88ad111..9df75792d2 100644 --- a/test/unit/test_function.py +++ b/test/unit/test_function.py @@ -371,27 +371,31 @@ class TestFunctionBackward(unittest.TestCase): def test_backward_single_call(self): N = 4 - x = Tensor.ones(N, N) - w1 = Tensor.ones(N, N, requires_grad=True) - w2 = Tensor.ones(N, N, requires_grad=True) + x = Tensor.arange(N*N).reshape(N, N).float() + w1 = Tensor.arange(N*N).reshape(N, N).float().requires_grad_() + w2 = Tensor.arange(N*N).reshape(N, N).float().requires_grad_() + xn, w1n, w2n = x.numpy(), w1.numpy(), w2.numpy() @function def f(t:Tensor, w1:Tensor, w2:Tensor): return (t@w1)@w2 f(x, w1, w2).sum().backward() assert w1.grad is not None and w2.grad is not None - np.testing.assert_allclose(w1.grad.numpy(), np.ones((N,N)) @ np.ones((N,N)) @ np.ones((N,N)), atol=1e-6) - np.testing.assert_allclose(w2.grad.numpy(), (np.ones((N,N)) @ np.ones((N,N))).T @ np.ones((N,N)), atol=1e-6) + np.testing.assert_allclose(w1.grad.numpy(), xn.T @ np.ones((N,N)) @ w2n.T, atol=1e-3) + np.testing.assert_allclose(w2.grad.numpy(), (xn @ w1n).T @ np.ones((N,N)), atol=1e-3) def test_backward_precompile_backward(self): N = 4 - x = Tensor.ones(N, N) - w1 = Tensor.ones(N, N, requires_grad=True) - w2 = Tensor.ones(N, N, requires_grad=True) - @function(precompile_backward=True) + x = Tensor.arange(N*N).reshape(N, N).float().contiguous() + w1 = Tensor.arange(N*N).reshape(N, N).float().requires_grad_().contiguous() + w2 = Tensor.arange(N*N).reshape(N, N).float().requires_grad_().contiguous() + Tensor.realize(x, w1, w2) + xn, w1n, w2n = x.numpy(), w1.numpy(), w2.numpy() + @function(precompile=True, precompile_backward=True) def f(t:Tensor, w1:Tensor, w2:Tensor): return (t@w1)@w2 - f(x, w1, w2).sum().backward() + loss = f(x, w1, w2).sum().backward() assert w1.grad is not None and w2.grad is not None - np.testing.assert_allclose(w1.grad.numpy(), np.ones((N,N)) @ np.ones((N,N)) @ np.ones((N,N)), atol=1e-6) - np.testing.assert_allclose(w2.grad.numpy(), (np.ones((N,N)) @ np.ones((N,N))).T @ np.ones((N,N)), atol=1e-6) + Tensor.realize(loss, w1.grad, w2.grad) + np.testing.assert_allclose(w1.grad.numpy(), xn.T @ np.ones((N,N)) @ w2n.T, atol=1e-3) + np.testing.assert_allclose(w2.grad.numpy(), (xn @ w1n).T @ np.ones((N,N)), atol=1e-3) if __name__ == '__main__': unittest.main() diff --git a/tinygrad/engine/allocations.py b/tinygrad/engine/allocations.py index 5a1fe9e794..ccdaec8612 100644 --- a/tinygrad/engine/allocations.py +++ b/tinygrad/engine/allocations.py @@ -107,7 +107,7 @@ def transform_precompiled_call(c:UOp) -> UOp|None: sink_srcs.append(target.after(target.store(elem))) fxn = UOp.sink(*sink_srcs) new_call = c.replace(src=(fxn, *input_buffers, *out_bufs), dtype=dtypes.void, tag=None) - return UOp.tuple(*[buf.after(new_call) for buf in out_bufs]) + return UOp.maketuple(*[buf.after(new_call) for buf in out_bufs]) out = _buffer_like(c) target = out.param_like(len(c.src)-1).shrink_to(c.shape) fxn = target.after(target.store(c.src[0])).sink() diff --git a/tinygrad/function.py b/tinygrad/function.py index d6be245a2a..ff9eab5366 100644 --- a/tinygrad/function.py +++ b/tinygrad/function.py @@ -43,7 +43,7 @@ class _function(Generic[ReturnType]): if isinstance(ret, Tensor): uret = ret.uop elif isinstance(ret, tuple) and all(isinstance(x, Tensor) for x in ret): - uret = UOp.tuple(*[x.uop for x in ret]) + uret = UOp.maketuple(*[x.uop for x in ret]) else: raise RuntimeError(f"function return type {type(ret)} not supported") diff --git a/tinygrad/gradient.py b/tinygrad/gradient.py index 06d935a8b1..59a1388330 100644 --- a/tinygrad/gradient.py +++ b/tinygrad/gradient.py @@ -29,7 +29,7 @@ def call_gradient(ctx:UOp, k:UOp) -> tuple[UOp|None, ...]: grad_uops.append(grads[p]) if len(grad_uops) == 0: return (None,) * (len(args) + 1) # build a single backward CALL returning a TUPLE of all gradients - bwd_body = UOp.tuple(*grad_uops) + bwd_body = UOp.maketuple(*grad_uops) bwd_call = bwd_body.call(*args, ctx, name=(k.arg.name or "")+"_backward", precompile=k.arg.precompile_backward) # extract each gradient via GETTUPLE ret: list[UOp|None] = [None] diff --git a/tinygrad/uop/ops.py b/tinygrad/uop/ops.py index 084ad8315a..c7d67d882c 100644 --- a/tinygrad/uop/ops.py +++ b/tinygrad/uop/ops.py @@ -410,7 +410,7 @@ class UOp(OpMixin, metaclass=UOpMetaClass): def sink(*srcs:UOp|None, **kwargs): # pylint: disable=no-self-argument return UOp(Ops.SINK, dtypes.void, tuple([x for x in srcs if x is not None]), **kwargs) - def tuple(*srcs:UOp): # pylint: disable=no-self-argument + def maketuple(*srcs:UOp): # pylint: disable=no-self-argument return UOp(Ops.TUPLE, dtypes.void, srcs) def gettuple(self, idx:int) -> UOp: in_tuple = self.src[0] if self.op is Ops.CALL else self