remove Tensor.from_uop (#16344)

and no device for const in Tensor init

test/backend/test_schedule.py

@@ -149,11 +149,6 @@ class TestSchedule(unittest.TestCase):
     run_linear(*check_schedule(z, 1, [x,y]))
     self.assertEqual(z.item(), 32)
 
-  def test_constants_can_store(self):
-    a = Tensor(2).contiguous()
-    run_linear(*check_schedule(a, 1))
-    np.testing.assert_equal(a.numpy(), 2)
-
   def test_allow_push_permutes(self):
     a = Tensor.randn(10,10,10).realize()
     b = Tensor.randn(10,10,1).realize()
@@ -712,7 +707,7 @@ class TestSchedule(unittest.TestCase):
     x = Tensor.empty(3,3,3,3)
     y = x.pad((-1,2,2,-1), mode="replicate")
     dx = y.sum().gradient(x)[0]
-    sched = check_schedule(dx, 1)
+    sched = check_schedule(dx, 0)
     run_linear(*sched)
     np.testing.assert_allclose(dx.numpy(), [[[[0.,3.,9.],[0,1.,3.],[0.,0.,0.]]]*3]*3)
 
@@ -1289,7 +1284,7 @@ class TestView(unittest.TestCase):
 class TestCopyFolding(unittest.TestCase):
   def test_const_copy_is_free(self):
     b = Tensor(1).to("CPU") * 4
-    run_linear(*check_schedule(b, 1, filter_sink=False))
+    run_linear(*check_schedule(b, 0, filter_sink=False))
     assert b.item() == 4
 
   def test_one_hot_with_copy(self):

test/null/test_schedule.py

@@ -766,12 +766,6 @@ class TestSchedule(unittest.TestCase):
     out = x + y
     check_schedule(out, 1)
 
-  def test_const_no_recompute(self):
-    x = Tensor(2) + Tensor(2)
-    y = Tensor(2) + Tensor(2)
-    out = x.contiguous() + y.contiguous()
-    check_schedule(out, 2, filter_sink=False)
-
   def test_reduce_shrink_child(self):
     a = Tensor.empty(100, 100)
     b = Tensor.empty(10,)
@@ -993,7 +987,7 @@ class TestSchedule(unittest.TestCase):
   def test_fuse_arange_pad_circular_mode_bw(self):
     x = Tensor.empty(1,1,5,5,5)
     out = x.pad((1,2,3,5,1,2), mode="circular")
-    g = out.sum().gradient(x)[0]
+    g = out.sum().gradient(x)[0].clone()
     linear, _ = check_schedule(g, 1)
     self.assertEqual(len([x for x in linear.src[0].src[0].backward_slice_with_self if x.op is Ops.REDUCE]), 0)
 

tinygrad/function.py

@@ -40,7 +40,7 @@ class _function(Generic[ReturnType]):
     params = get_state_dict((args, kwargs), tensor_type=(Tensor, UOp)).values()
 
     # deduplicate input_uops, keeping the first occurrence index for each unique uop
-    call_uops: list[UOp] = dedup([(t.uop if isinstance(t, Tensor) else t) for t in params])
+    call_uops: list[UOp] = dedup([u for t in params if not ((u:=(t.uop if isinstance(t, Tensor) else t)).base.op is Ops.CONST and u.device is None)])
 
     # disable realize/schedule while this is running
     # run it and do surgery later

tinygrad/gradient.py

@@ -22,14 +22,17 @@ def _compact_params(body:UOp, all_args:tuple[UOp, ...]) -> tuple[UOp, tuple[UOp,
 def call_gradient(ctx:UOp, k:UOp, needed:set[int]) -> tuple[UOp|None, ...]:
   fxn, args = k.src[0], k.src[1:]
   if k.arg.grad_fxn is not None:
+    # put const on a device, also TODO why do we still have NOOP...
+    def on_dev(g, i): return g.clone(device=args[i].device if k.op is Ops.CALL else k.device) if g.device is None else g
     if ctx.op is Ops.TUPLE:
-      real = [g for g in ctx.src if g.op is not Ops.NOOP]
+      real = [on_dev(g, i) for i,g in enumerate(ctx.src) if g.op is not Ops.NOOP]
       return (None,) + (k.arg.grad_fxn(*real, call=k) if len(real) > 1 else k.arg.grad_fxn(real[0], k))
-    return (None,) + k.arg.grad_fxn(ctx, k)
+    return (None,) + k.arg.grad_fxn(on_dev(ctx, 0), k)
   assert fxn.op is Ops.TUPLE, f"expected TUPLE body for gradient, got {fxn.op}"
   params = {x.arg:x for x in fxn.toposort(enter_calls=False) if x.op == Ops.PARAM}
   grad_args = ctx.src
-  root_grad = UOp(Ops.TUPLE, src=tuple(UOp(Ops.NOOP) if g.op is Ops.NOOP else g.param_like(len(args)+i) for i,g in enumerate(grad_args)))
+  root_grad = UOp(Ops.TUPLE, src=tuple(UOp(Ops.NOOP) if g.op is Ops.NOOP else
+    g if g.base.op is Ops.CONST and g.device is None else g.param_like(len(args)+i) for i,g in enumerate(grad_args)))
   grads = compute_gradient(fxn, root_grad, set(params.values()))
   # for precompiled calls, substitute forward outputs with params so intermediates aren't recomputed
   fwd_subs = {src: src.param_like(len(args)+len(grad_args)+i) for i, src in enumerate(fxn.src)} if k.arg.precompile else {}

tinygrad/tensor.py

@@ -109,11 +109,11 @@ class Tensor(OpMixin):
     if isinstance(data, UOp):
       assert _dtype is None or _dtype==data.dtype or data.dtype==dtypes.weakint, f"dtype mismatch: {_dtype} vs {data.dtype}"
       # if data is dtype.weakint that means that this is a symbolic int and we need to lower it to something we can make a Tensor out of
-      if data.dtype == dtypes.weakint: data = Tensor.from_uop(data, device=_device).uop
+      if data.dtype == dtypes.weakint: data = _index_to_concrete_int(data)
     elif data is None:
-      data = UOp.const(_dtype or dtypes.default_float, 0, _device)
+      data = UOp.const(_dtype or dtypes.default_float, 0)
     elif isinstance(data, get_args(ConstType)):
-      data = UOp.const(_dtype or dtypes.from_py(data), data, _device)
+      data = UOp.const(_dtype or dtypes.from_py(data), data)
     elif isinstance(data, bytes): data = _frompy(data, _dtype or dtypes.uint8, _device)
     elif isinstance(data, (list, tuple)):
       if _dtype is None:
@@ -125,7 +125,7 @@ class Tensor(OpMixin):
       import numpy as np
       assert isinstance(data, np.ndarray), f"expected np.ndarray, got {data}"
       if data.shape == ():
-        data = UOp.const(_dtype or _from_np_dtype(data.dtype), data.item(), _device)
+        data = UOp.const(_dtype or _from_np_dtype(data.dtype), data.item())
       else:
         data = _fromnp(data.astype(npdtype) if _dtype is not None and (npdtype:=_to_np_dtype(_dtype)) is not None else data)
     elif isinstance(data, pathlib.Path):
@@ -260,7 +260,7 @@ class Tensor(OpMixin):
     # broadcast x (shape only, dtype must match)
     if self.shape != x.shape: x = x._broadcast_to(self.shape)
     if self.shape != x.shape: raise RuntimeError(f"assign shape mismatch {self.shape} != {x.shape}")
-    if not is_disk and x.uop.device is not None and self.device != x.device:
+    if not is_disk and x.uop.device is not None and self.device is not None and self.device != x.device:
       raise RuntimeError(f"assign device mismatch {self.device} != {x.device}")
     if not is_disk and self.dtype != x.dtype: raise RuntimeError(f"assign dtype mismatch {self.dtype} != {x.dtype}")
     if isinstance(self.device, tuple) and self.uop.axis != x.uop.axis: raise RuntimeError(f"multi axis mismatch {self.uop.axis} != {x.uop.axis}")
@@ -463,20 +463,6 @@ class Tensor(OpMixin):
 
     return data[:16].contiguous()
 
-  @staticmethod
-  def from_uop(y:UOp, **kwargs) -> Tensor:
-    # TODO: remove this and stay in weakint
-    if y.dtype == dtypes.weakint: y = _index_to_concrete_int(y)
-    if y.op is Ops.BIND:
-      var, val = y.unbind()
-      _device = canonicalize_device(kwargs.get("device"))
-      const = UOp.const(var.dtype, val, _device, ())
-      return Tensor(y.replace(src=(var.replace(src=const.src), const)), **kwargs)
-    if y.op is Ops.CONST: return Tensor(y.arg, **kwargs)
-    if y.op is Ops.MUL: return Tensor.from_uop(y.src[0]) * Tensor.from_uop(y.src[1])
-    if y.op is Ops.ADD: return Tensor.from_uop(y.src[0]) + Tensor.from_uop(y.src[1])
-    raise RuntimeError(f"unhandled UOp {y}")
-
   # ***** creation entrypoint *****
 
   @staticmethod

tinygrad/uop/render.py

@@ -108,10 +108,10 @@ pm_pyrender_extra = PatternMatcher([
   (UPat(Ops.CMOD, name="x"), lambda ctx,x: f"{ctx[x.src[0]]}.alu(Ops.CMOD, {ctx[x.src[1]]})"),
   # NOTE: only match CONSTs without UNIQUE (len(src)==1), unique_const needs explicit rendering
   (UPat(set(syms.keys())-{Ops.SUB, Ops.CMPNE, Ops.CDIV, Ops.CMOD}, src=(UPat(Ops.CONST, src=(UPat(Ops.DEVICE),), name="y"), UPat(name="z")),
-        name="x"), lambda ctx,x,y,z: strip_binary_parens(x, str(y.arg), ctx[z], lambda a,b: f"({a}{syms[x.op]}{b})")),
+        name="x"), lambda ctx,x,y,z: strip_binary_parens(x, str(y.arg), ctx[z], lambda a,b: f"({a}{syms[x.op]}{b})") if y.device==z.device else None),
   # NOTE: sub doesn't work cause it's written as add/mul
   (UPat(set(syms.keys())-{Ops.SUB, Ops.CDIV, Ops.CMOD}, src=(UPat(name="y"), UPat(Ops.CONST, src=(UPat(Ops.DEVICE),), name="z")), name="x"),
-    lambda ctx,x,y,z: strip_binary_parens(x, ctx[y], str(z.arg), lambda a,b: f"({a}{syms[x.op]}{b})")),
+    lambda ctx,x,y,z: strip_binary_parens(x, ctx[y], str(z.arg), lambda a,b: f"({a}{syms[x.op]}{b})") if y.device==z.device else None),
   (UPat(set(syms.keys())-{Ops.SUB, Ops.CDIV, Ops.CMOD}, name="x"), lambda ctx,x:
     strip_binary_parens(x, ctx[x.src[0]], ctx[x.src[1]], lambda a,b: f"({a}{syms[x.op]}{b})")),
   (UPat(sugar, src=(), name="x"), lambda x: f"UOp.{x.op.name.lower()}("+', '.join(([f'arg={repr(x.arg)}'] if x.arg is not None else []))+")"),