more deviceless const prerequisites [pr] (#16256)

* more deviceless const prerequisites [pr]

* remove that

* arange.contiguous -> arange.clone in tests

arange will become deviceless const soon, update tests where it needs to be a buffer

test/backend/test_multitensor.py

@@ -1174,23 +1174,23 @@ class TestMultiBufferView(unittest.TestCase):
     self._check(ref, a, lambda t: t[3])
 
   def test_shrink_2d(self):
-    ref = Tensor.arange(6*4).reshape(6, 4).contiguous().realize()
-    a = Tensor.arange(6*4).reshape(6, 4).contiguous().shard(devices_2, axis=1).realize()
+    ref = Tensor.arange(6*4).reshape(6, 4).clone().realize()
+    a = Tensor.arange(6*4).reshape(6, 4).clone().shard(devices_2, axis=1).realize()
     self._check(ref, a, lambda t: t.shrink(((1, 4), None)))
 
   def test_reshape_then_shrink(self):
-    ref = Tensor.arange(8*6).reshape(8, 6).contiguous().realize()
-    a = Tensor.arange(8*6).reshape(8, 6).contiguous().shard(devices_2, axis=1).realize()
+    ref = Tensor.arange(8*6).reshape(8, 6).clone().realize()
+    a = Tensor.arange(8*6).reshape(8, 6).clone().shard(devices_2, axis=1).realize()
     self._check(ref, a, lambda t: t.reshape(4, 2, 6)[1])
 
   def test_chained_shrink(self):
-    ref = Tensor.arange(10*8).reshape(10, 8).contiguous().realize()
-    a = Tensor.arange(10*8).reshape(10, 8).contiguous().shard(devices_2, axis=1).realize()
+    ref = Tensor.arange(10*8).reshape(10, 8).clone().realize()
+    a = Tensor.arange(10*8).reshape(10, 8).clone().shard(devices_2, axis=1).realize()
     self._check(ref, a, lambda t: t.shrink(((2, 8), None)).shrink(((1, 4), None)))
 
   def test_4_devices(self):
-    ref = Tensor.arange(8*12).reshape(8, 12).contiguous().realize()
-    a = Tensor.arange(8*12).reshape(8, 12).contiguous().shard(devices_4, axis=1).realize()
+    ref = Tensor.arange(8*12).reshape(8, 12).clone().realize()
+    a = Tensor.arange(8*12).reshape(8, 12).clone().shard(devices_4, axis=1).realize()
     out = a[5].contiguous()
     linear, var_vals = out.linear_with_vars()
     if all(hasattr(Device[d].allocator, "_offset") for d in out.device):

test/backend/test_schedule.py

@@ -966,7 +966,7 @@ class TestSchedule(unittest.TestCase):
   def test_arange_index_contiguous(self):
     Tensor.manual_seed(0)
     x = Tensor.randn(5, 2).realize()
-    a = Tensor.arange(10).contiguous()
+    a = Tensor.arange(10).clone()
     out = (x + a[2]).sum()
     run_linear(*check_schedule(out, 2))
     np.testing.assert_allclose(out.numpy(), (x.numpy()+np.arange(10)[2]).sum(), atol=1e-5, rtol=1e-6)
@@ -982,7 +982,7 @@ class TestSchedule(unittest.TestCase):
   def test_user_contiguous(self):
     Tensor.manual_seed(0)
     x = Tensor.randn(5, 2).realize()
-    a = (Tensor.arange(10)+1).contiguous()
+    a = (Tensor.arange(10)+1).clone()
     out = (x + a[2]).sum()
     run_linear(*check_schedule(out, 2))
     np.testing.assert_allclose(out.numpy(), (x.numpy()+(np.arange(10)+1)[2]).sum(), atol=1e-5, rtol=1e-6)
@@ -1008,7 +1008,7 @@ class TestSchedule(unittest.TestCase):
   def test_fuse_assign_contiguous(self):
     x = Tensor.zeros(4, 4, dtype=dtypes.int).contiguous().realize()
     a = Tensor.arange(8).reshape(4, 2)
-    run_linear(*check_schedule(x.shrink((None, (0, 2))).assign(a.contiguous()), 2))
+    run_linear(*check_schedule(x.shrink((None, (0, 2))).assign(a.clone()), 2))
     np.testing.assert_equal(x.numpy(), [[0, 1, 0, 0], [2, 3, 0, 0], [4, 5, 0, 0], [6, 7, 0, 0]])
 
   def test_assign_non_contiguous_alt(self): self.test_assign_non_contiguous(alt=True)
@@ -1053,7 +1053,7 @@ class TestSchedule(unittest.TestCase):
 
   def test_no_extra_contiguous_on_setitem_assign_back(self):
     # pattern: contiguous copy, advanced setitem, assign back (e.g. torch backend _view_write)
-    base = Tensor.arange(16).reshape(4, 4).contiguous()
+    base = Tensor.arange(16).reshape(4, 4).clone()
     flat_base = base.reshape(16).contiguous()
     idx = Tensor([1,2,5,6], dtype=dtypes.int32)
     flat_base[idx] = Tensor([99,99,99,99])
@@ -1253,7 +1253,7 @@ class TestView(unittest.TestCase):
   # x collapses along with its children
   def test_parent_view_collapses(self):
     a = Tensor([1, 2])
-    b = Tensor.arange(3).contiguous()
+    b = Tensor.arange(3).clone()
     bv = b.pad(((0, 2),))[-2:]
     # this becomes a late a*0
     late_mul = a*bv
@@ -1270,7 +1270,7 @@ class TestView(unittest.TestCase):
   # as long as one child realizes, x does not collapse
   def test_parent_multiple_children_no_collapse(self):
     a = Tensor([1, 2])
-    b = Tensor.arange(3).contiguous()
+    b = Tensor.arange(3).clone()
     bv = b.pad(((0, 2),))[-2:]
     late_mul = a*bv
     other_child = b+2

test/backend/test_tensor.py

@@ -36,16 +36,6 @@ class TestTinygrad(unittest.TestCase):
     self.assertTrue(t.uop.has_buffer_identity())
     np.testing.assert_equal(t.numpy(), 2.0)
 
-  def test_to_deviceless_const(self):
-    t = Tensor(UOp.const(dtypes.float, 2.0))
-    self.assertIs(t.to(f"{Device.DEFAULT}:1"), t)
-    self.assertIs(t.to_(f"{Device.DEFAULT}:1"), t)
-
-  def test_shard_deviceless_const(self):
-    t = Tensor(UOp.const(dtypes.float, 2.0))
-    self.assertIs(t.shard((f"{Device.DEFAULT}:0", f"{Device.DEFAULT}:1")), t)
-    self.assertIs(t.shard_((f"{Device.DEFAULT}:0", f"{Device.DEFAULT}:1")), t)
-
   def test_plus_equals(self):
     a = Tensor.randn(10,10)
     b = Tensor.randn(10,10)

test/null/test_schedule.py

@@ -1216,17 +1216,17 @@ class TestFusionOp(unittest.TestCase):
 class TestBufferView(unittest.TestCase):
   def test_shrink_contiguous_is_buffer_view(self):
     # simple 1D shrink of a realized buffer should be BUFFER_VIEW, not a copy kernel
-    a = Tensor.arange(100).contiguous().realize()
+    a = Tensor.arange(100).clone().realize()
     b = a.shrink(((10, 50),)).contiguous()
     run_linear(*check_schedule(b, 0))
 
   def test_shrink_2d_contiguous_is_buffer_view(self):
-    a = Tensor.arange(100).reshape(10,10).contiguous().realize()
+    a = Tensor.arange(100).reshape(10,10).clone().realize()
     b = a.shrink(((1, 5),None)).contiguous()
     run_linear(*check_schedule(b, 0))
 
   def test_chained_shrink_is_buffer_view(self):
-    a = Tensor.arange(1000).contiguous().realize()
+    a = Tensor.arange(1000).clone().realize()
     b = a.shrink(((200, 800),)).shrink(((0, 300),)).reshape((30, 10)).shrink(((20, 25), (0, 10))).contiguous()
     run_linear(*check_schedule(b, 0))
 
@@ -1234,96 +1234,96 @@ class TestBufferView(unittest.TestCase):
     # indexing a non-shard axis of a realized sharded tensor should be BUFFER_VIEW on each device, not copy kernels
     # this is the flat_llama pattern: weight[layer_idx] where weight is (n_layers, out, dim) sharded on axis=1
     devices = ("NULL:1", "NULL:2")
-    a = Tensor.arange(8*4*10).reshape(8, 4, 10).contiguous().shard(devices, axis=1).realize()
+    a = Tensor.arange(8*4*10).reshape(8, 4, 10).clone().shard(devices, axis=1).realize()
     run_linear(*check_schedule(a[3].contiguous(), 0))
 
   def test_shrink_2d_non_shard_axis_multi(self):
     devices = ("NULL:1", "NULL:2")
-    a = Tensor.arange(6*4).reshape(6, 4).contiguous().shard(devices, axis=1).realize()
+    a = Tensor.arange(6*4).reshape(6, 4).clone().shard(devices, axis=1).realize()
     run_linear(*check_schedule(a.shrink(((1, 4), None)).contiguous(), 0))
 
   def test_shrink_shard_axis_0_multi(self):
     # shrinking a middle dim is not contiguous per shard, so this needs copy kernels
     devices = ("NULL:1", "NULL:2")
-    a = Tensor.arange(4*6*2).reshape(4, 6, 2).contiguous().shard(devices, axis=0).realize()
+    a = Tensor.arange(4*6*2).reshape(4, 6, 2).clone().shard(devices, axis=0).realize()
     run_linear(*check_schedule(a.shrink((None, (2, 5), None)).contiguous(), 2))
 
   def test_reshape_then_shrink_multi(self):
     devices = ("NULL:1", "NULL:2")
-    a = Tensor.arange(8*6).reshape(8, 6).contiguous().shard(devices, axis=1).realize()
+    a = Tensor.arange(8*6).reshape(8, 6).clone().shard(devices, axis=1).realize()
     run_linear(*check_schedule(a.reshape(4, 2, 6)[1].contiguous(), 0))
 
   def test_permute_then_shrink_multi(self):
     # permute makes per-shard view non-contiguous, needs copy kernels
     devices = ("NULL:1", "NULL:2")
-    a = Tensor.arange(4*6*2).reshape(4, 6, 2).contiguous().shard(devices, axis=1).realize()
+    a = Tensor.arange(4*6*2).reshape(4, 6, 2).clone().shard(devices, axis=1).realize()
     run_linear(*check_schedule(a.permute(1, 0, 2).shrink(((0, 6), (1, 3), None)).contiguous(), 2))
 
   def test_multi_buffer_view_4_devices(self):
     devices = tuple(f"NULL:{i}" for i in range(4))
-    a = Tensor.arange(8*12).reshape(8, 12).contiguous().shard(devices, axis=1).realize()
+    a = Tensor.arange(8*12).reshape(8, 12).clone().shard(devices, axis=1).realize()
     run_linear(*check_schedule(a[5].contiguous(), 0))
 
   def test_chained_shrink_multi(self):
     devices = ("NULL:1", "NULL:2")
-    a = Tensor.arange(10*8).reshape(10, 8).contiguous().shard(devices, axis=1).realize()
+    a = Tensor.arange(10*8).reshape(10, 8).clone().shard(devices, axis=1).realize()
     run_linear(*check_schedule(a.shrink(((2, 8), None)).shrink(((1, 4), None)).contiguous(), 0))
 
   # negative tests: these should NOT become BUFFER_VIEW (non-contiguous per shard)
   def test_expand_multi_not_buffer_view(self):
     devices = ("NULL:1", "NULL:2")
-    a = Tensor.arange(4*2).reshape(4, 1, 2).contiguous().shard(devices, axis=2).realize()
+    a = Tensor.arange(4*2).reshape(4, 1, 2).clone().shard(devices, axis=2).realize()
     run_linear(*check_schedule(a.expand(4, 3, 2).contiguous(), 2))
 
   def test_pad_multi_not_buffer_view(self):
     devices = ("NULL:1", "NULL:2")
-    a = Tensor.arange(4*2).reshape(4, 2).contiguous().shard(devices, axis=1).realize()
+    a = Tensor.arange(4*2).reshape(4, 2).clone().shard(devices, axis=1).realize()
     run_linear(*check_schedule(a.pad(((1, 1), (0, 0))).contiguous(), 2))
 
   def test_flip_multi_not_buffer_view(self):
     devices = ("NULL:1", "NULL:2")
-    a = Tensor.arange(4*2).reshape(4, 2).contiguous().shard(devices, axis=1).realize()
+    a = Tensor.arange(4*2).reshape(4, 2).clone().shard(devices, axis=1).realize()
     run_linear(*check_schedule(a.flip(0).contiguous(), 2))
 
   def test_replicated_reshape_is_buffer_view(self):
     devices = ("NULL:1", "NULL:2")
-    a = Tensor.arange(24).contiguous().to(devices).realize()
+    a = Tensor.arange(24).clone().to(devices).realize()
     run_linear(*check_schedule(a.reshape(4, 6).contiguous(), 0))
 
   def test_replicated_shrink_is_buffer_view(self):
     # DP pattern: replicated weight[layer_idx]
     devices = ("NULL:1", "NULL:2")
-    a = Tensor.arange(8*10).reshape(8, 10).contiguous().to(devices).realize()
+    a = Tensor.arange(8*10).reshape(8, 10).clone().to(devices).realize()
     run_linear(*check_schedule(a[3].contiguous(), 0))
 
   def test_replicated_chained_mops_is_buffer_view(self):
     devices = ("NULL:1", "NULL:2")
-    a = Tensor.arange(100).contiguous().to(devices).realize()
+    a = Tensor.arange(100).clone().to(devices).realize()
     run_linear(*check_schedule(a.reshape(10, 10).shrink(((2, 7), None)).contiguous(), 0))
 
   def test_replicated_shard_none_is_buffer_view(self):
     devices = ("NULL:1", "NULL:2")
-    a = Tensor.arange(24).contiguous().shard(devices, axis=None).realize()
+    a = Tensor.arange(24).clone().shard(devices, axis=None).realize()
     run_linear(*check_schedule(a.reshape(4, 6).contiguous(), 0))
 
   def test_replicated_4_devices_is_buffer_view(self):
     devices = tuple(f"NULL:{i}" for i in range(4))
-    a = Tensor.arange(8*10).reshape(8, 10).contiguous().to(devices).realize()
+    a = Tensor.arange(8*10).reshape(8, 10).clone().to(devices).realize()
     run_linear(*check_schedule(a[3].contiguous(), 0))
 
   def test_replicated_expand_not_buffer_view(self):
     devices = ("NULL:1", "NULL:2")
-    a = Tensor.arange(12).reshape(4, 1, 3).contiguous().to(devices).realize()
+    a = Tensor.arange(12).reshape(4, 1, 3).clone().to(devices).realize()
     run_linear(*check_schedule(a.expand(4, 3, 3).contiguous(), 2))
 
   def test_replicated_permute_not_buffer_view(self):
     devices = ("NULL:1", "NULL:2")
-    a = Tensor.arange(24).reshape(4, 6).contiguous().to(devices).realize()
+    a = Tensor.arange(24).reshape(4, 6).clone().to(devices).realize()
     run_linear(*check_schedule(a.permute(1, 0).contiguous(), 2))
 
   def test_replicated_flip_not_buffer_view(self):
     devices = ("NULL:1", "NULL:2")
-    a = Tensor.arange(24).reshape(4, 6).contiguous().to(devices).realize()
+    a = Tensor.arange(24).reshape(4, 6).clone().to(devices).realize()
     run_linear(*check_schedule(a.flip(0).contiguous(), 2))
 
 class TestInvalidTensor(unittest.TestCase):

tinygrad/mixin/__init__.py

@@ -935,7 +935,8 @@ class OpMixin(ElementwiseMixin, ReduceMixin):
     print(t.gather(1, Tensor([[0, 0], [1, 0]])).numpy())
     ```
     """
-    if index.device != self.device: raise RuntimeError(f"expected index and self on the same device, {index.device=}, {self.device=}")
+    if index.device is not None and self.device is not None and index.device != self.device:
+      raise RuntimeError(f"expected index and self on the same device, {index.device=}, {self.device=}")
     if index.ndim != self.ndim: raise RuntimeError(f"self.ndim must equal index.ndim, {self.ndim=}, {index.ndim=}")
     dim = self._resolve_dim(dim)
     assert all(s >= i for d,(s,i) in enumerate(zip(self.shape, index.shape)) if d != dim), "requires self.shape[d] >= index.shape[d] for all d != dim"
@@ -975,8 +976,10 @@ class OpMixin(ElementwiseMixin, ReduceMixin):
     return x.cast(self.dtype)
 
   def _pre_scatter(self, dim:int, index:Self, src:Self) -> tuple[Self, Self]:
-    if index.device != self.device: raise RuntimeError(f"expected index and self on the same device, {index.device=}, {self.device=}")
-    if src.device != self.device: raise RuntimeError(f"expected src and self on the same device, {src.device=}, {self.device=}")
+    if index.device is not None and self.device is not None and index.device != self.device:
+      raise RuntimeError(f"expected index and self on the same device, {index.device=}, {self.device=}")
+    if src.device is not None and self.device is not None and src.device != self.device:
+      raise RuntimeError(f"expected src and self on the same device, {src.device=}, {self.device=}")
     dim = self._resolve_dim(dim)
     assert index.ndim == self.ndim == src.ndim, f"self.ndim, index.ndim and src.ndim must all equal, {self.ndim=} {index.ndim=} {src.ndim=}"
     assert all((d == dim or self_ >= index_) and src_ >= index_ for d,(self_,index_,src_) in enumerate(zip(self.shape, index.shape, src.shape))), \
@@ -1393,7 +1396,8 @@ class OpMixin(ElementwiseMixin, ReduceMixin):
     ```
     """
     assert 0.0 <= label_smoothing <= 1.0, "label_smoothing must be in [0.0, 1.0]"
-    if Y.device != self.device: raise RuntimeError(f"expected Y and self on the same device, {Y.device=}, {self.device=}")
+    if Y.device is not None and self.device is not None and Y.device != self.device:
+      raise RuntimeError(f"expected Y and self on the same device, {Y.device=}, {self.device=}")
     log_probs = self.log_softmax()
     loss_mask = Y.ne(ignore_index) if ignore_index != -1 else Y.ones_like(dtype=dtypes.bool)
     y = Y.unsqueeze(-1)._one_hot_along_dim(self.shape[-1], dim=-1) * loss_mask.unsqueeze(-1)

tinygrad/tensor.py

@@ -887,7 +887,8 @@ class Tensor(OpMixin):
       match index:
         case Tensor():
           if not dtypes.is_int(index.dtype): raise IndexError(f"index dtype {index.dtype} is not supported")
-          if index.device != self.device: raise RuntimeError(f"expected index and self on the same device, {index.device=}, {self.device=}")
+          if index.device is not None and self.device is not None and index.device != self.device:
+            raise RuntimeError(f"expected index and self on the same device, {index.device=}, {self.device=}")
           assert isinstance(size, int), "size must be an int"
           index = (index < 0).where(index+size, index)  # treat negative index values
         case list() | tuple():