From 188d7ec15e6c4c557d27bae9b1e608bc46b603b4 Mon Sep 17 00:00:00 2001 From: chenyu Date: Tue, 19 May 2026 21:29:27 -0400 Subject: [PATCH] clone can take device (#16271) useful to materialize const on a specific device --- test/backend/test_tensor.py | 5 +++++ tinygrad/tensor.py | 12 +++++++----- 2 files changed, 12 insertions(+), 5 deletions(-) diff --git a/test/backend/test_tensor.py b/test/backend/test_tensor.py index 3249a6a3fe..d1229e3cd4 100644 --- a/test/backend/test_tensor.py +++ b/test/backend/test_tensor.py @@ -755,6 +755,11 @@ class TestZeroShapeTensor(unittest.TestCase): 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")) diff --git a/tinygrad/tensor.py b/tinygrad/tensor.py index 909e308599..1cd7961d97 100644 --- a/tinygrad/tensor.py +++ b/tinygrad/tensor.py @@ -226,7 +226,7 @@ class Tensor(OpMixin): 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 @@ class Tensor(OpMixin): 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: """