From 66306b532196047ecf26eb75f20ebbbad2ec3abe Mon Sep 17 00:00:00 2001 From: FICTURE7 Date: Sun, 2 Feb 2025 15:50:34 +0400 Subject: [PATCH] Fix disk tensor assignment (#8855) * Add test for disk tensor assignment failure * Fix disk tensor assignment --------- Co-authored-by: qazal <77887910+Qazalin@users.noreply.github.com> --- test/test_assign.py | 5 +++++ tinygrad/tensor.py | 2 +- 2 files changed, 6 insertions(+), 1 deletion(-) diff --git a/test/test_assign.py b/test/test_assign.py index 1cb0b99d47..bc22bf12e3 100644 --- a/test/test_assign.py +++ b/test/test_assign.py @@ -3,6 +3,7 @@ import unittest import numpy as np from tinygrad import dtypes, Tensor, TinyJit, GlobalCounters, Variable from tinygrad.device import is_dtype_supported +from tinygrad.helpers import temp N = 200 # has to be bigger than the cache to fail @@ -388,5 +389,9 @@ class TestAssign(unittest.TestCase): assert oba1 is None and oba2 is None np.testing.assert_allclose(a.numpy(), np.arange(N*N,dtype=np.int32).reshape((N,N))) + def test_disk_assignment(self): + a = Tensor.empty(5, device=f"disk:{temp('disk_assignment')}").assign(Tensor.ones(5)).numpy() + np.testing.assert_equal(a, np.ones(5)) + if __name__ == "__main__": unittest.main() diff --git a/tinygrad/tensor.py b/tinygrad/tensor.py index f32e5139cd..8d048eb431 100644 --- a/tinygrad/tensor.py +++ b/tinygrad/tensor.py @@ -279,7 +279,7 @@ class Tensor(SimpleMathTrait): # 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="CLANG", dtype=self.dtype) - self.contiguous().realize().lazydata.base.realized.copyin(x._data()) + self.contiguous().realize().lazydata.base.realized.ensure_allocated().copyin(x._data()) return self if x.__class__ is not Tensor: x = Tensor(x, device=self.device, dtype=self.dtype) if self.lazydata is x.lazydata: return self # a self assign is a NOOP