more prerequisuite test changed for deviceless const (#16264)

This commit is contained in:
chenyu
2026-05-19 15:43:45 -04:00
committed by GitHub
parent aa1e59ab97
commit 890b731b1e
5 changed files with 22 additions and 22 deletions
+2 -2
View File
@@ -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()
+5 -5
View File
@@ -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
View File
@@ -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
View File
@@ -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)
+2 -2
View File
@@ -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)