clone can take device (tinygrad#16271)

useful to materialize const on a specific device

Author: chenyuxyz

test/backend/test_tensor.py

@@ -755,6 +755,11 @@ def test_clone_with_grad(self):
     assert b.grad is not None
     np.testing.assert_allclose(a.grad.numpy(), b.grad.numpy())
 
+  def test_clone_deviceless_const_to_cpu(self):
+    t = Tensor(UOp.const(dtypes.float, 2.0)).clone(device="CPU")
+    self.assertEqual(t.device, "CPU")
+    np.testing.assert_equal(t.numpy(), 2.0)
+
   def test_reduce_default(self):
     np.testing.assert_equal(Tensor([]).max().numpy(), -float("inf"))
     np.testing.assert_equal(Tensor([]).min().numpy(), float("inf"))

tinygrad/tensor.py

@@ -226,7 +226,7 @@ def callify(self, *lst:Tensor) -> Tensor:
   def linear_with_vars(self, *lst:Tensor) -> tuple[UOp, dict[str, int]]:
     """Creates the LINEAR UOp needed to realize these Tensor(s), with Variables."""
     for x in (self,)+lst:
-      if x.uop.device is None: x.replace(Tensor.empty(*x.shape, dtype=x.dtype, device=Device.DEFAULT).assign(x))
+      if x.uop.device is None: x.replace(x.clone(device=Device.DEFAULT))
     big_sink, becomes_map = transform_to_call(UOp.sink(*[x.uop for x in (self,)+lst]))
     _apply_map_to_tensors(becomes_map, name="buffers")
     return create_linear_with_vars(big_sink)
@@ -353,13 +353,15 @@ def numpy(self) -> 'numpy.ndarray':
     if 0 in self.shape: return np.empty(self.shape, dtype=_to_np_dtype(self.dtype.base))
     return self._buffer().numpy().reshape(self.shape)
 
-  def clone(self) -> Tensor:
+  def clone(self, device:str|tuple[str, ...]|None=None) -> Tensor:
     """
     Creates a clone of this tensor allocating a separate buffer for the data.
+    If `device` is specified, the clone is placed on that device.
     """
-    ret = self.empty_like()
-    if self.grad is not None: ret.grad = self.grad.clone()
-    return ret.assign(self)
+    device = device or self.device
+    ret = self.empty_like(device=device)
+    if self.grad is not None: ret.grad = self.grad.clone(device=device)
+    return ret.assign(self.to(device))
 
   def to(self, device:str|tuple[str, ...]|None) -> Tensor:
     """