forked from tinygrad/tinygrad
more prerequisuite test changed for deviceless const (#16264)
This commit is contained in:
@@ -122,7 +122,7 @@ class TestReduceOpsConstFolding(unittest.TestCase):
|
||||
class TestMultiConstFolding(unittest.TestCase):
|
||||
def test_multi_const_folding_literal(self):
|
||||
ds = tuple(f"{Device.DEFAULT}:{i}" for i in range(4))
|
||||
t = Tensor.arange(16).float().to(ds).realize()
|
||||
t = Tensor.arange(16).float().clone().to(ds).realize()
|
||||
|
||||
# non const folding case creates one ast on each shard
|
||||
_check_ast_count(4, t + 1)
|
||||
@@ -147,7 +147,7 @@ class TestMultiConstFolding(unittest.TestCase):
|
||||
|
||||
def test_multi_const_folding_tensor(self):
|
||||
ds = tuple(f"{Device.DEFAULT}:{i}" for i in range(4))
|
||||
t = Tensor.arange(16).float().to(ds).realize()
|
||||
t = Tensor.arange(16).float().clone().to(ds).realize()
|
||||
zero = Tensor.zeros(16).to(ds).realize()
|
||||
one = Tensor.ones(16).to(ds).realize()
|
||||
|
||||
|
||||
@@ -122,7 +122,7 @@ class TestMultiTensor(unittest.TestCase):
|
||||
_ = Tensor(X.uop, dtype=dtypes.float)
|
||||
|
||||
def test_sharded_arange(self):
|
||||
sharded_arange = Tensor.arange(1000).shard(devices_2, 0)
|
||||
sharded_arange = Tensor.arange(1000).clone().shard(devices_2, 0)
|
||||
sharded_arange.realize()
|
||||
np.testing.assert_equal(sharded_arange.numpy(), np.arange(1000))
|
||||
|
||||
@@ -235,7 +235,7 @@ class TestMultiTensor(unittest.TestCase):
|
||||
for ring in (0, 2):
|
||||
GlobalCounters.reset()
|
||||
with Context(RING=ring, SCACHE=0):
|
||||
t = Tensor.arange(32).contiguous().shard(devices_4, 0).to(Device.DEFAULT)
|
||||
t = Tensor.arange(32).clone().shard(devices_4, 0).to(Device.DEFAULT)
|
||||
t.realize()
|
||||
kernel_counts[ring] = GlobalCounters.kernel_count
|
||||
self.assertEqual(t.device, Device.DEFAULT)
|
||||
@@ -557,7 +557,7 @@ class TestMultiTensor(unittest.TestCase):
|
||||
def test_multi_tensor_jit_graph_assign_updates_each_shard(self):
|
||||
@TinyJit
|
||||
def jf(out: Tensor) -> Tensor:
|
||||
tmp = (Tensor.arange(4, dtype=dtypes.float).shard(devices_2, 0) + 1).contiguous().realize()
|
||||
tmp = (Tensor.arange(4, dtype=dtypes.float).clone().shard(devices_2, 0) + 1).contiguous().realize()
|
||||
out.assign((tmp + 1).contiguous()).realize()
|
||||
return out
|
||||
|
||||
@@ -839,7 +839,7 @@ class TestMultiTensor(unittest.TestCase):
|
||||
|
||||
def test_clone(self):
|
||||
for axis in (None, 0):
|
||||
t = Tensor.arange(16).reshape(4, 4).shard(devices_2, axis=axis).contiguous().realize()
|
||||
t = Tensor.arange(16).reshape(4, 4).clone().shard(devices_2, axis=axis).contiguous().realize()
|
||||
t_clone = t.clone().realize()
|
||||
self.assertEqual(t_clone.device, t.device)
|
||||
self.assertEqual(t_clone.uop.axis, axis)
|
||||
@@ -1170,7 +1170,7 @@ class TestMultiBufferView(unittest.TestCase):
|
||||
@unittest.skip("flaky on LLVM")
|
||||
def test_shrink_non_shard_axis(self):
|
||||
ref = Tensor.arange(8*4*10).reshape(8, 4, 10).contiguous().realize()
|
||||
a = Tensor.arange(8*4*10).reshape(8, 4, 10).contiguous().shard(devices_2, axis=1).realize()
|
||||
a = Tensor.arange(8*4*10).reshape(8, 4, 10).clone().shard(devices_2, axis=1).realize()
|
||||
self._check(ref, a, lambda t: t[3])
|
||||
|
||||
def test_shrink_2d(self):
|
||||
|
||||
+3
-3
@@ -40,7 +40,7 @@ class TestExample(unittest.TestCase):
|
||||
|
||||
@multidevice_test
|
||||
def test_example_readme(self, device):
|
||||
x = Tensor.eye(3, device=device)
|
||||
x = Tensor.eye(3).clone().to(device)
|
||||
y = Tensor([[2.0,0,-2.0]], device=device)
|
||||
z = y.matmul(x).sum()
|
||||
z.backward()
|
||||
@@ -59,8 +59,8 @@ class TestExample(unittest.TestCase):
|
||||
print(f"WARNING: {device} test isn't running")
|
||||
return
|
||||
|
||||
x = Tensor.eye(8, device=device)
|
||||
y = Tensor.eye(8, device=device)
|
||||
x = Tensor.eye(8).clone().to(device)
|
||||
y = Tensor.eye(8).clone().to(device)
|
||||
z = y.matmul(x).sum()
|
||||
z.backward()
|
||||
|
||||
|
||||
+10
-10
@@ -230,7 +230,7 @@ class TestCallSchedule(unittest.TestCase):
|
||||
@function(precompile=True)
|
||||
def f(x:Tensor) -> Tensor: return x * 2 + 1
|
||||
sz = UOp.variable("sz", 1, 16)
|
||||
a = Tensor.arange(16*4).reshape(16, 4).float()[:sz.bind(5)]
|
||||
a = Tensor.arange(16*4).reshape(16, 4).float().clone()[:sz.bind(5)]
|
||||
out = f(a)
|
||||
# result shape should have the symbolic dim, not the max
|
||||
self.assertIsInstance(out.shape[0], UOp)
|
||||
@@ -240,7 +240,7 @@ class TestCallSchedule(unittest.TestCase):
|
||||
@function(precompile=True)
|
||||
def f(x:Tensor) -> Tensor: return x + 1
|
||||
devs = ("CPU:0", "CPU:1")
|
||||
a = Tensor.arange(8).reshape(4, 2).float().shard(devs, axis=0)
|
||||
a = Tensor.arange(8).reshape(4, 2).float().clone().shard(devs, axis=0)
|
||||
out = f(a) + 2
|
||||
np.testing.assert_allclose(out.numpy(), np.arange(8, dtype=np.float32).reshape(4, 2) + 3)
|
||||
|
||||
@@ -251,7 +251,7 @@ class TestCallMultiSharded(unittest.TestCase):
|
||||
devs = ("CPU:0", "CPU:1")
|
||||
@function
|
||||
def f(x:Tensor): return (x + 1, x * 2)
|
||||
a = Tensor.arange(8).reshape(4, 2).float().shard(devs, axis=0)
|
||||
a = Tensor.arange(8).reshape(4, 2).float().clone().shard(devs, axis=0)
|
||||
t1, t2 = f(a)
|
||||
ref = np.arange(8, dtype=np.float32).reshape(4, 2)
|
||||
np.testing.assert_allclose(t1.numpy(), ref + 1)
|
||||
@@ -262,7 +262,7 @@ class TestCallMultiSharded(unittest.TestCase):
|
||||
devs = ("CPU:0", "CPU:1")
|
||||
@function(precompile=True)
|
||||
def f(x:Tensor): return (x + 1, x * 2)
|
||||
a = Tensor.arange(8).reshape(4, 2).float().shard(devs, axis=0)
|
||||
a = Tensor.arange(8).reshape(4, 2).float().clone().shard(devs, axis=0)
|
||||
t1, t2 = f(a)
|
||||
ref = np.arange(8, dtype=np.float32).reshape(4, 2)
|
||||
np.testing.assert_allclose(t1.numpy(), ref + 1)
|
||||
@@ -273,7 +273,7 @@ class TestCallMultiSharded(unittest.TestCase):
|
||||
devs = ("CPU:0", "CPU:1")
|
||||
@function
|
||||
def f(x:Tensor): return (x.sum(axis=0), x.sum(axis=1))
|
||||
a = Tensor.arange(8).reshape(4, 2).float().shard(devs, axis=0)
|
||||
a = Tensor.arange(8).reshape(4, 2).float().clone().shard(devs, axis=0)
|
||||
t1, t2 = f(a)
|
||||
ref = np.arange(8, dtype=np.float32).reshape(4, 2)
|
||||
np.testing.assert_allclose(t1.numpy(), ref.sum(axis=0))
|
||||
@@ -284,8 +284,8 @@ class TestCallMultiSharded(unittest.TestCase):
|
||||
devs = ("CPU:0", "CPU:1")
|
||||
@function
|
||||
def f(x:Tensor, y:Tensor): return (x + y, x * y)
|
||||
a = Tensor.arange(8).reshape(4, 2).float().shard(devs, axis=0)
|
||||
b = Tensor.arange(8).reshape(4, 2).float().shard(devs, axis=0) + 1
|
||||
a = Tensor.arange(8).reshape(4, 2).float().clone().shard(devs, axis=0)
|
||||
b = Tensor.arange(8).reshape(4, 2).float().clone().shard(devs, axis=0) + 1
|
||||
t1, t2 = f(a, b)
|
||||
ref_a = np.arange(8, dtype=np.float32).reshape(4, 2)
|
||||
ref_b = ref_a + 1
|
||||
@@ -297,7 +297,7 @@ class TestCallMultiSharded(unittest.TestCase):
|
||||
devs = ("CPU:0", "CPU:1")
|
||||
@function
|
||||
def f(x:Tensor): return (x + 1, x * 2)
|
||||
a = Tensor.arange(8).reshape(4, 2).float().shard(devs, axis=0)
|
||||
a = Tensor.arange(8).reshape(4, 2).float().clone().shard(devs, axis=0)
|
||||
t1, t2 = f(a)
|
||||
out = (t1 + t2).sum()
|
||||
ref = np.arange(8, dtype=np.float32).reshape(4, 2)
|
||||
@@ -308,8 +308,8 @@ class TestCallMultiSharded(unittest.TestCase):
|
||||
devs = ("CPU:0", "CPU:1")
|
||||
@function
|
||||
def f(x:Tensor, y:Tensor): return (x + 1, y + 2)
|
||||
a = Tensor.arange(8).reshape(4, 2).float().shard(devs, axis=0)
|
||||
b = Tensor.arange(8).reshape(4, 2).float().shard(devs, axis=1)
|
||||
a = Tensor.arange(8).reshape(4, 2).float().clone().shard(devs, axis=0)
|
||||
b = Tensor.arange(8).reshape(4, 2).float().clone().shard(devs, axis=1)
|
||||
t1, t2 = f(a, b)
|
||||
ref_a = np.arange(8, dtype=np.float32).reshape(4, 2)
|
||||
ref_b = np.arange(8, dtype=np.float32).reshape(4, 2)
|
||||
|
||||
@@ -306,7 +306,7 @@ class TestDiskTensor(TempDirTestCase):
|
||||
dt[::2] = Tensor([10, 20, 30])
|
||||
|
||||
def test_advanced_setitem_not_supported(self):
|
||||
dt = Tensor.arange(12).reshape(3, 4).to(f"disk:{self.tmp('dt_advanced_setitem')}")
|
||||
dt = Tensor.arange(12).reshape(3, 4).clone().to(f"disk:{self.tmp('dt_advanced_setitem')}")
|
||||
with self.assertRaises(RuntimeError, msg="advanced setitem is not supported for DISK tensors"):
|
||||
dt[Tensor([0, 2]), Tensor([1, 3])] = 99
|
||||
|
||||
@@ -549,7 +549,7 @@ class TestDiskTensorMovement(TempDirTestCase):
|
||||
def setUp(self):
|
||||
super().setUp()
|
||||
self.fn = pathlib.Path(self.tmp("custom_disk_range"))
|
||||
Tensor.arange(100, dtype=dtypes.uint8).to(f"disk:{str(self.fn)}").realize()
|
||||
Tensor.arange(100, dtype=dtypes.uint8).clone().to(f"disk:{str(self.fn)}").realize()
|
||||
|
||||
def test_simple_read(self):
|
||||
t = Tensor(self.fn)
|
||||
|
||||
Reference in New Issue
Block a user