maketuple

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()

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()

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")
 

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]

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