From 86a204d22ae209556b99eb1a9d520da1c14509a2 Mon Sep 17 00:00:00 2001 From: chenyu Date: Thu, 29 Jan 2026 21:26:58 -0500 Subject: [PATCH] allow Tensor setitem input to be list/tuple (#14432) matches assign, and generally matches numpy --- test/unit/test_assign.py | 5 +++++ tinygrad/tensor.py | 11 +++++------ 2 files changed, 10 insertions(+), 6 deletions(-) diff --git a/test/unit/test_assign.py b/test/unit/test_assign.py index 24c2ce60e4..076f2b14f3 100644 --- a/test/unit/test_assign.py +++ b/test/unit/test_assign.py @@ -456,6 +456,11 @@ class TestAssign(unittest.TestCase): assign.realize() np.testing.assert_allclose(a.numpy(), [2., 2., 2., 2., 1., 1., 1., 1.]) + def test_setitem_list(self): + a = Tensor.zeros(8).contiguous().realize() + a[2:5] = [1, 2, 3] + np.testing.assert_allclose(a.numpy(), [0., 0., 1., 2., 3., 0., 0., 0.]) + @unittest.skip("don't use output buffer, and mismatch dtype no longer supported") def test_cast_assignment(self): a = Tensor(np.arange(N*N, dtype=np.float32)).reshape(N,N) diff --git a/tinygrad/tensor.py b/tinygrad/tensor.py index 2e664f5484..430a7f04a2 100644 --- a/tinygrad/tensor.py +++ b/tinygrad/tensor.py @@ -275,13 +275,13 @@ class Tensor(OpMixin): self.uop = x.uop return self - def assign(self, x) -> Tensor: + def assign(self, x:Tensor|PyConst|list|tuple) -> Tensor: # TODO: this is a hack for writing to DISK. remove with working assign if isinstance(self.device, str) and self.device.startswith("DISK"): - if x.__class__ is not Tensor: x = Tensor(x, device="CPU", dtype=self.dtype) + if not isinstance(x, Tensor): x = Tensor(x, device="CPU", dtype=self.dtype) self._buffer().copyin(x._data()) return self - if x.__class__ is not Tensor: x = Tensor(x, device=self.device, dtype=self.dtype) + if not isinstance(x, Tensor): x = Tensor(x, device=self.device, dtype=self.dtype) if self.uop is x.uop: return self # a self assign is a NOOP # NOTE: we allow cross device assign # broadcast x @@ -1268,13 +1268,12 @@ class Tensor(OpMixin): """ return self._getitem(indices) - def __setitem__(self, indices, v:Tensor|PyConst) -> None: + def __setitem__(self, indices, v:Tensor|PyConst|list|tuple) -> None: if isinstance(self.device, str) and self.device.startswith("DISK"): self.realize()._getitem(indices).assign(v) return # NOTE: check that setitem target is valid first - if isinstance(v, get_args(PyConst)): v = Tensor(v, device=self.device, dtype=self.dtype) - if not isinstance(v, Tensor): raise TypeError(f"can't set a {type(v).__name__} to a Tensor") + if not isinstance(v, Tensor): v = Tensor(v, device=self.device, dtype=self.dtype) if self.requires_grad or v.requires_grad: raise NotImplementedError("setitem with requires_grad is not supported") self.realize() if not self.uop.is_contiguous(): raise RuntimeError("setitem target needs to be contiguous")