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>
This commit is contained in:
FICTURE7
2025-02-02 15:50:34 +04:00
committed by GitHub
parent 6e523e4d17
commit 66306b5321
2 changed files with 6 additions and 1 deletions

View File

@@ -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()