forked from tinygrad/tinygrad
test update for arange and eye (#16394)
these will need explicit clone to make a buffer
This commit is contained in:
@@ -155,7 +155,7 @@ class TestCustomKernel(unittest.TestCase):
|
||||
self.assertTrue((ref == tst).all().item())
|
||||
|
||||
def test_eye(self):
|
||||
ref = Tensor.eye(1024).contiguous().realize()
|
||||
ref = Tensor.eye(1024).clone().realize()
|
||||
tst = Tensor.empty_like(ref)
|
||||
tst = tst.custom_kernel(fxn=custom_eye_kernel)[0]
|
||||
self.assertTrue((ref == tst).all().item())
|
||||
@@ -335,7 +335,7 @@ class TestCustomKernel(unittest.TestCase):
|
||||
assert all(x == expected for x in result), f"expected all {expected}, got {result}"
|
||||
|
||||
def test_custom_kernel_sched(self, use_custom=False):
|
||||
x = Tensor.arange(32).reshape(8, 4).realize()
|
||||
x = Tensor.arange(32).reshape(8, 4).clone().realize()
|
||||
y = Tensor.empty_like(x)
|
||||
y = Tensor.custom_kernel(y, x, fxn=custom_add_one_kernel)[0]
|
||||
if use_custom:
|
||||
@@ -352,7 +352,7 @@ class TestCustomKernel(unittest.TestCase):
|
||||
|
||||
@unittest.expectedFailure
|
||||
def test_sliced_buffer_function(self):
|
||||
x = Tensor.arange(32).reshape(8, 4).realize()
|
||||
x = Tensor.arange(32).reshape(8, 4).clone().realize()
|
||||
from tinygrad import function
|
||||
@function(precompile=True)
|
||||
def run(x:Tensor) -> Tensor:
|
||||
|
||||
@@ -49,7 +49,7 @@ class TestJit(unittest.TestCase):
|
||||
y = (x + 1).contiguous().realize()
|
||||
z = x.shrink(((st, st + N),)).contiguous().realize()
|
||||
return y, z
|
||||
x = Tensor.arange(2*N).contiguous().realize()
|
||||
x = Tensor.arange(2*N).clone().realize()
|
||||
for _ in range(3): y, z = f(x, Variable("a", 0, N).bind(0))
|
||||
self.assertEqual(y.shape, (2*N,))
|
||||
self.assertEqual(z.shape, (N,))
|
||||
@@ -92,7 +92,7 @@ class TestJit(unittest.TestCase):
|
||||
@TinyJit
|
||||
def f(x): return (x[2:5].contiguous() + 1).realize()
|
||||
for i in range(5):
|
||||
x = (Tensor.arange(10).float() + i * 10).contiguous().realize()
|
||||
x = (Tensor.arange(10).float() + i * 10).clone().realize()
|
||||
np.testing.assert_allclose(f(x).numpy(), x.numpy()[2:5] + 1)
|
||||
|
||||
def test_jit_multiple_outputs(self):
|
||||
|
||||
@@ -187,7 +187,7 @@ class TestMultiTensor(unittest.TestCase):
|
||||
np.testing.assert_allclose(O.numpy(), X.numpy()[0:2]*W.numpy()[0:2] < 2)
|
||||
|
||||
def test_shrink_on_shard_axis(self):
|
||||
X = Tensor.arange(4*4).reshape(4,4).realize()
|
||||
X = Tensor.arange(4*4).reshape(4,4).clone().realize()
|
||||
X_np = X.numpy()
|
||||
X.shard_(devices_2, 0)
|
||||
# only shrink on the device that owns the shard, this is enabled by the mselect simplifier
|
||||
@@ -293,7 +293,7 @@ class TestMultiTensor(unittest.TestCase):
|
||||
@TinyJit
|
||||
def f(x): return (x+1).contiguous().sum()
|
||||
for _ in range(5):
|
||||
tt = Tensor.arange(0, 4).contiguous().realize().shard((d1,d2), 0).realize()
|
||||
tt = Tensor.arange(0, 4).clone().realize().shard((d1,d2), 0).realize()
|
||||
out = f(tt)
|
||||
assert out.item() == 1+2+3+4
|
||||
|
||||
@@ -309,7 +309,7 @@ class TestMultiTensor(unittest.TestCase):
|
||||
@TinyJit
|
||||
def f(x): return (x.shard((d1,d2), 0)+1).contiguous().sum()
|
||||
for _ in range(5):
|
||||
tt = Tensor.arange(0, 4).contiguous().realize()
|
||||
tt = Tensor.arange(0, 4).clone().realize()
|
||||
out = f(tt)
|
||||
assert out.item() == 1+2+3+4
|
||||
|
||||
@@ -865,7 +865,7 @@ class TestMultiTensor(unittest.TestCase):
|
||||
@unittest.skip("RANGEIFY doesn't support multi const folding")
|
||||
def test_multi_const_folding(self):
|
||||
with Context(TRACK_MATCH_STATS=0):
|
||||
a = Tensor.arange(3).realize()
|
||||
a = Tensor.arange(3).clone().realize()
|
||||
zeros = Tensor.zeros(3).realize()
|
||||
b = a.to(devices_2)*zeros.to(devices_2)
|
||||
sched = b.schedule_linear().src
|
||||
@@ -904,7 +904,7 @@ class TestShrinkMultiTensorShardedAxis(unittest.TestCase):
|
||||
|
||||
# shrink a multitensor on sharded axis
|
||||
def test_shrink_bad_args(self):
|
||||
t = Tensor.arange(64).reshape(8, 8).contiguous().realize()
|
||||
t = Tensor.arange(64).reshape(8, 8).clone().realize()
|
||||
t.shard_([f"{Device.DEFAULT}:{i}" for i in range(4)], axis=0)
|
||||
|
||||
with self.assertRaises(AssertionError):
|
||||
@@ -927,7 +927,7 @@ class TestShrinkMultiTensorShardedAxis(unittest.TestCase):
|
||||
@given(strat.sampled_from([dtypes.float, dtypes.int, dtypes.int64, dtypes.int16]))
|
||||
def test_ops(self, dtype):
|
||||
if dtype not in Device[Device.DEFAULT].renderer.supported_dtypes(): return
|
||||
t = Tensor.arange(64).reshape(8, 8).contiguous().realize()
|
||||
t = Tensor.arange(64).reshape(8, 8).clone().realize()
|
||||
t.shard_([f"{Device.DEFAULT}:{i}" for i in range(4)], axis=0)
|
||||
for i in range(4):
|
||||
print(f"{i=}")
|
||||
@@ -971,7 +971,7 @@ class TestShrinkMultiTensorShardedAxis(unittest.TestCase):
|
||||
np.testing.assert_allclose(a.flip(-1).numpy(), b.flip(-1).numpy(), rtol=1e-7, atol=1e-3)
|
||||
|
||||
def test_add_two_partitions(self):
|
||||
t = Tensor.arange(64).reshape(8, 8).contiguous().realize()
|
||||
t = Tensor.arange(64).reshape(8, 8).clone().realize()
|
||||
t.shard_([f"{Device.DEFAULT}:{i}" for i in range(4)], axis=0)
|
||||
|
||||
a = t.shrink(((2, 4), None))
|
||||
@@ -988,7 +988,7 @@ class TestShrinkMultiTensorShardedAxis(unittest.TestCase):
|
||||
|
||||
def test_add_different_tensors(self):
|
||||
devices = [f"{Device.DEFAULT}:{i}" for i in range(4)]
|
||||
x = Tensor.arange(64).reshape(8, 8).contiguous().realize().shard(devices, axis=0)
|
||||
x = Tensor.arange(64).reshape(8, 8).clone().realize().shard(devices, axis=0)
|
||||
|
||||
to_add = []
|
||||
for i in range(len(devices)):
|
||||
@@ -1098,7 +1098,7 @@ class TestBatchNorm(unittest.TestCase):
|
||||
@given(strat.sampled_from((False, True)))
|
||||
def test_batchnorm(self, is_training):
|
||||
devices = [f"{Device.DEFAULT}:{i}" for i in range(4)]
|
||||
x = Tensor.arange(4096).reshape(8, 8, 8, 8).contiguous().realize().shard(devices, axis=0)
|
||||
x = Tensor.arange(4096).reshape(8, 8, 8, 8).clone().realize().shard(devices, axis=0)
|
||||
|
||||
with Tensor.train(is_training):
|
||||
bns = []
|
||||
@@ -1184,7 +1184,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()
|
||||
ref = Tensor.arange(8*4*10).reshape(8, 4, 10).clone().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])
|
||||
|
||||
@@ -1296,7 +1296,7 @@ class TestMultiSetitem(unittest.TestCase):
|
||||
@needs_second_gpu
|
||||
def setUp(self): pass
|
||||
|
||||
def _t(self, axis): return Tensor.arange(16).contiguous().realize().shard(self.device, axis=axis)
|
||||
def _t(self, axis): return Tensor.arange(16).clone().realize().shard(self.device, axis=axis)
|
||||
|
||||
def test_setitem_scalar_axis0(self):
|
||||
t = self._t(0)
|
||||
|
||||
@@ -67,7 +67,7 @@ class TestPickle(unittest.TestCase):
|
||||
|
||||
# NOTE: currently Buffer exists on the uop, not tensor
|
||||
def test_pickle_buffer_uop(self):
|
||||
t = Tensor.arange(4).realize()
|
||||
t = Tensor.arange(4).clone().realize()
|
||||
a = t.uop
|
||||
assert a.is_realized
|
||||
self.assertIsNotNone(buffer:=a.base.realized)
|
||||
@@ -95,7 +95,7 @@ class TestPickle(unittest.TestCase):
|
||||
np.testing.assert_equal(vt2.numpy(), 20)
|
||||
|
||||
def test_pickle_buffer_view(self):
|
||||
t = Tensor.arange(10, device="CPU").contiguous().realize()
|
||||
t = Tensor.arange(10).clone(device="CPU").realize()
|
||||
vt = t[3:5].contiguous().realize()
|
||||
assert hasattr(vt.uop.buffer, 'base')
|
||||
ref_value = vt.tolist()
|
||||
|
||||
@@ -221,8 +221,8 @@ class TestSchedule(unittest.TestCase):
|
||||
np.testing.assert_allclose(out.numpy(), (x.numpy() - x.numpy().max(keepdims=True)).max())
|
||||
|
||||
def test_example_matmul_contig(self):
|
||||
x = Tensor.eye(64).contiguous().realize()
|
||||
y = Tensor.eye(64).contiguous().realize()
|
||||
x = Tensor.eye(64).clone().realize()
|
||||
y = Tensor.eye(64).clone().realize()
|
||||
z = y.matmul(x).sum()
|
||||
z.backward()
|
||||
out = x.grad.contiguous()
|
||||
@@ -826,7 +826,7 @@ class TestSchedule(unittest.TestCase):
|
||||
self._test_fusion([(32, 32)], lambda a:a-a.sum(1), 2)
|
||||
|
||||
def test_cast_padded_view(self):
|
||||
a = Tensor.arange(4).reshape(1, 4)
|
||||
a = Tensor.arange(4).reshape(1, 4).clone().realize()
|
||||
casted_view = a.pad(((0, 1), (0, 0))).cast(dtypes.float)
|
||||
casted_view.realize()
|
||||
self.assertEqual(casted_view.uop.base.realized.size, 8)
|
||||
@@ -836,7 +836,7 @@ class TestSchedule(unittest.TestCase):
|
||||
|
||||
# NOTE: we only reorder CAST if it's an EXPAND
|
||||
def test_cast_after_shrink(self):
|
||||
a = Tensor.arange(4).reshape(1, 4)
|
||||
a = Tensor.arange(4).reshape(1, 4).clone().realize()
|
||||
casted_view = a.shrink(((0, 1), (0, 2))).cast(dtypes.float)
|
||||
casted_view.realize()
|
||||
self.assertEqual(casted_view.uop.base.realized.size, 2)
|
||||
@@ -991,7 +991,7 @@ class TestSchedule(unittest.TestCase):
|
||||
|
||||
def test_assign_non_contiguous_alt(self): self.test_assign_non_contiguous(alt=True)
|
||||
def test_assign_non_contiguous(self, alt=False):
|
||||
x = (Tensor.arange(16)-100).reshape(4,4).contiguous().realize()
|
||||
x = (Tensor.arange(16)-100).reshape(4,4).clone().realize()
|
||||
xref = x.numpy()
|
||||
if alt:
|
||||
y = Tensor.randint(2, 4).contiguous().realize()
|
||||
@@ -1007,7 +1007,7 @@ class TestSchedule(unittest.TestCase):
|
||||
np.testing.assert_equal(tst.numpy(), a.numpy())
|
||||
|
||||
def test_setitem_sched(self, mop=lambda x:x, expected_kcount=1):
|
||||
a = Tensor.arange(16, device="CPU").reshape(4, 4).contiguous().realize()
|
||||
a = Tensor.arange(16).reshape(4, 4).clone(device="CPU").realize()
|
||||
a2 = mop(a)
|
||||
expected = (a+a2).tolist()
|
||||
a.assign(a+a2)
|
||||
@@ -1021,7 +1021,7 @@ class TestSchedule(unittest.TestCase):
|
||||
|
||||
def test_setitem_const_fused(self):
|
||||
# https://github.com/tinygrad/tinygrad/issues/10690
|
||||
a = Tensor.arange(16).contiguous().realize()
|
||||
a = Tensor.arange(16).clone().realize()
|
||||
GlobalCounters.reset()
|
||||
a[4] = 3
|
||||
self.assertEqual(GlobalCounters.kernel_count, 0)
|
||||
@@ -1278,7 +1278,7 @@ class TestCopyFolding(unittest.TestCase):
|
||||
self.assertEqual(x.item(), 2.0)
|
||||
|
||||
def test_late_const_copy_folding(self):
|
||||
a = Tensor.arange(3).realize()
|
||||
a = Tensor.arange(3).clone().realize()
|
||||
zeros = Tensor.zeros(3, buffer=False).realize()
|
||||
b = (a*zeros).to("CPU") + 1
|
||||
run_linear(*check_schedule(b, 1, filter_sink=False))
|
||||
@@ -1353,14 +1353,14 @@ class TestCopyFolding(unittest.TestCase):
|
||||
self.assertListEqual(b.tolist(), [[0, 2], [1, 3]])
|
||||
|
||||
def test_permute_on_disk(self):
|
||||
with open(temp('dt_arange_4_permute'), "wb") as f: f.write(Tensor.arange(4).realize().uop.base.buffer.as_memoryview())
|
||||
with open(temp('dt_arange_4_permute'), "wb") as f: f.write(Tensor.arange(4).clone().realize().uop.base.buffer.as_memoryview())
|
||||
a = Tensor.empty(4, dtype=dtypes.int32, device=f"disk:{temp('dt_arange_4_permute')}")
|
||||
b = a.reshape(2, 2).permute(1, 0).to("CPU")
|
||||
b.realize()
|
||||
self.assertListEqual(b.tolist(), [[0, 2], [1, 3]])
|
||||
|
||||
def test_permute_on_disk_contiguous(self):
|
||||
with open(temp('dt_arange_4_permute_contig'), "wb") as f: f.write(Tensor.arange(4).realize().uop.base.buffer.as_memoryview())
|
||||
with open(temp('dt_arange_4_permute_contig'), "wb") as f: f.write(Tensor.arange(4).clone().realize().uop.base.buffer.as_memoryview())
|
||||
a = Tensor.empty(4, dtype=dtypes.int32, device=f"disk:{temp('dt_arange_4_permute_contig')}")
|
||||
b = a.reshape(2, 2).permute(1, 0).contiguous().to("CPU")
|
||||
b.realize()
|
||||
@@ -1374,7 +1374,7 @@ class TestCopyFolding(unittest.TestCase):
|
||||
|
||||
# NOTE: disk permute must come after COPY
|
||||
def test_permute_after_shrink_on_disk(self):
|
||||
with open(temp('dt_arange_5_permute'), "wb") as f: f.write(Tensor.arange(5).realize().uop.base.buffer.as_memoryview())
|
||||
with open(temp('dt_arange_5_permute'), "wb") as f: f.write(Tensor.arange(5).clone().realize().uop.base.buffer.as_memoryview())
|
||||
a = Tensor.empty(5, dtype=dtypes.int32, device=f"disk:{temp('dt_arange_5_permute')}")
|
||||
b = a.shrink(((0, 4),)).reshape(2, 2).permute(1, 0).to("CPU")
|
||||
b.realize()
|
||||
|
||||
@@ -32,13 +32,13 @@ class TestSetitem(unittest.TestCase):
|
||||
self.assertListEqual(t.tolist(), [0, 1, 11, 3, 11, 5, 6, 7, 8, 9])
|
||||
|
||||
def test_setitem_inplace_mul(self):
|
||||
t = Tensor.arange(10).realize()
|
||||
t = Tensor.arange(10).clone().realize()
|
||||
t[:3] *= 10
|
||||
self.assertListEqual(t.tolist(), [0, 10, 20, 3, 4, 5, 6, 7, 8, 9])
|
||||
|
||||
def test_setitem_fancy_on_unrealized_view(self):
|
||||
# fancy indexing setitem on unrealized SHRINK view (triggered infinite loop in graph_rewrite)
|
||||
base = Tensor.arange(20, dtype=dtypes.float).reshape(4, 5)
|
||||
base = Tensor.arange(20, dtype=dtypes.float).reshape(4, 5).clone().realize()
|
||||
sub = base[1:3]
|
||||
flat = sub.reshape(sub.numel()).contiguous()
|
||||
idx = Tensor([0, 3, 7, 9])
|
||||
@@ -229,7 +229,7 @@ class TestSetitem(unittest.TestCase):
|
||||
np.testing.assert_equal(t.numpy(), n)
|
||||
|
||||
def test_setitem_swap_rows(self):
|
||||
t = Tensor.arange(6, dtype=dtypes.float).reshape(3, 2).contiguous().realize()
|
||||
t = Tensor.arange(6, dtype=dtypes.float).reshape(3, 2).clone().realize()
|
||||
tmp = t[0]
|
||||
t[0] = t[1]
|
||||
t[2] = tmp
|
||||
@@ -237,7 +237,7 @@ class TestSetitem(unittest.TestCase):
|
||||
np.testing.assert_allclose(t.numpy(), [[2, 3], [2, 3], [2, 3]])
|
||||
|
||||
# eager version
|
||||
t = Tensor.arange(6, dtype=dtypes.float).reshape(3, 2).contiguous().realize()
|
||||
t = Tensor.arange(6, dtype=dtypes.float).reshape(3, 2).clone().realize()
|
||||
tmp = t[0].realize()
|
||||
t[0] = t[1].realize()
|
||||
t[2] = tmp.realize()
|
||||
@@ -269,8 +269,8 @@ class TestSetitem(unittest.TestCase):
|
||||
def test_cross_assign_independence(self):
|
||||
# when assigning to two tensors using computations from both,
|
||||
# both assigns should see the OLD values of both tensors
|
||||
a = Tensor.arange(4, dtype=dtypes.float).contiguous().realize()
|
||||
b = Tensor.arange(4, 8, dtype=dtypes.float).contiguous().realize()
|
||||
a = Tensor.arange(4, dtype=dtypes.float).clone().realize()
|
||||
b = Tensor.arange(4, 8, dtype=dtypes.float).clone().realize()
|
||||
new_a = a + b # [4, 6, 8, 10]
|
||||
new_b = a * 2 # [0, 2, 4, 6] -- should use OLD a
|
||||
a.assign(new_a)
|
||||
@@ -283,8 +283,8 @@ class TestSetitem(unittest.TestCase):
|
||||
np.testing.assert_allclose(b.numpy(), [8, 12, 16, 20])
|
||||
|
||||
# eager version
|
||||
a = Tensor.arange(4, dtype=dtypes.float).contiguous().realize()
|
||||
b = Tensor.arange(4, 8, dtype=dtypes.float).contiguous().realize()
|
||||
a = Tensor.arange(4, dtype=dtypes.float).clone().realize()
|
||||
b = Tensor.arange(4, 8, dtype=dtypes.float).clone().realize()
|
||||
new_a = (a + b).realize()
|
||||
new_b = (a * 2).realize()
|
||||
a.assign(new_a).realize()
|
||||
|
||||
@@ -36,7 +36,7 @@ class TestSubBuffer(unittest.TestCase):
|
||||
assert len(mv) == 5
|
||||
|
||||
def test_subbuffer_used(self):
|
||||
t = Tensor.arange(0, 10, dtype=dtypes.uint8).realize()
|
||||
t = Tensor.arange(0, 10, dtype=dtypes.uint8).clone().realize()
|
||||
vt = t[2:4].realize()
|
||||
out = (vt + 100).tolist()
|
||||
assert out == [102, 103]
|
||||
@@ -44,7 +44,7 @@ class TestSubBuffer(unittest.TestCase):
|
||||
@needs_second_gpu
|
||||
@unittest.skipIf(Device.DEFAULT not in {"CUDA", "NV", "AMD"} or DEV.interface.startswith("MOCK"), "only NV, AMD, CUDA")
|
||||
def test_subbuffer_transfer(self):
|
||||
t = Tensor.arange(0, 10, dtype=dtypes.uint8).realize()
|
||||
t = Tensor.arange(0, 10, dtype=dtypes.uint8).clone().realize()
|
||||
vt = t[2:5].contiguous().realize()
|
||||
out = vt.to(f"{Device.DEFAULT}:1").realize().tolist()
|
||||
assert out == [2, 3, 4]
|
||||
|
||||
@@ -551,7 +551,7 @@ class TestTinygrad(unittest.TestCase):
|
||||
Tensor.zeros(2, 2).realize()
|
||||
|
||||
def test_shrink(self):
|
||||
t = Tensor.arange(32).contiguous().realize()
|
||||
t = Tensor.arange(32).clone().realize()
|
||||
self.assertListEqual(t[16:20].tolist(), [16,17,18,19])
|
||||
self.assertListEqual(t.shrink_to(16).tolist(), list(range(16)))
|
||||
t = t.reshape(4, 8).contiguous().realize()
|
||||
|
||||
Vendored
+1
-1
@@ -29,7 +29,7 @@ def gradient_test():
|
||||
z = y.matmul(x).sum()
|
||||
z.backward()
|
||||
def realized_eye():
|
||||
Tensor.eye(3).realize()
|
||||
Tensor.eye(3).clone().realize()
|
||||
def realized_list():
|
||||
Tensor([[2.0,0,-2.0]]).realize()
|
||||
def kernel_matmul():
|
||||
|
||||
@@ -134,14 +134,14 @@ class TestBitcastConstFolding(unittest.TestCase):
|
||||
# folds advance indexing into basic indexing
|
||||
class TestIndexingConstFolding(unittest.TestCase):
|
||||
def test_scalar_index(self):
|
||||
t = Tensor.arange(16).float().reshape(1,1,4,4).realize()
|
||||
t = Tensor.arange(16).float().reshape(1,1,4,4).clone().realize()
|
||||
_check_ast_count(1, t[:,:,Tensor(1),:])
|
||||
_check_ast_count(1, t[:,:,Tensor(1)+2,:])
|
||||
_check_ast_count(1, t[:,:,Tensor(1),Tensor(0)])
|
||||
|
||||
def test_const_tensor_index(self):
|
||||
# TODO: these can be 0, implement const tensor folded indexing
|
||||
t = Tensor.arange(16).float().reshape(1,1,4,4).realize()
|
||||
t = Tensor.arange(16).float().reshape(1,1,4,4).clone().realize()
|
||||
_check_ast_count(1, t[:,:,Tensor.ones(2,1,dtype=dtypes.int),:])
|
||||
_check_ast_count(1, t[:,:,Tensor.ones(1,2,dtype=dtypes.int)+2,:])
|
||||
_check_ast_count(1, t[:,:,Tensor.ones(1,1,dtype=dtypes.int),Tensor.zeros(2,1,2,dtype=dtypes.int)])
|
||||
|
||||
+23
-23
@@ -281,15 +281,15 @@ class TestAssign(unittest.TestCase):
|
||||
np.testing.assert_equal(t.numpy(), [[100, 104, 108, 112], [101, 105, 109, 113], [102, 106, 110, 114], [103, 107, 111, 115]])
|
||||
|
||||
def test_assign_contiguous(self):
|
||||
b = Tensor.arange(16).reshape(4,4).contiguous().realize()
|
||||
a = (Tensor.arange(16).reshape(4,4).contiguous().realize() + 1)
|
||||
b = Tensor.arange(16).reshape(4,4).clone().realize()
|
||||
a = (Tensor.arange(16).reshape(4,4).clone().realize() + 1)
|
||||
GlobalCounters.reset()
|
||||
b.assign(a.contiguous()).realize()
|
||||
self.assertEqual(GlobalCounters.kernel_count, 2)
|
||||
|
||||
def test_assign_contiguous_permute(self):
|
||||
b = Tensor.arange(16).reshape(4,4).contiguous().realize()
|
||||
a = (Tensor.arange(16).reshape(4,4).contiguous().realize() + 1).permute((1,0))
|
||||
b = Tensor.arange(16).reshape(4,4).clone().realize()
|
||||
a = (Tensor.arange(16).reshape(4,4).clone().realize() + 1).permute((1,0))
|
||||
GlobalCounters.reset()
|
||||
b.assign(a.contiguous()).realize()
|
||||
self.assertEqual(GlobalCounters.kernel_count, 2)
|
||||
@@ -325,29 +325,29 @@ class TestAssign(unittest.TestCase):
|
||||
np.testing.assert_allclose(a.numpy(), np.arange(N*N).reshape((N,N)) + np.arange(N*N).reshape((N,N)).transpose(1,0))
|
||||
|
||||
def test_post_permuted_assignment_alt(self):
|
||||
a = Tensor.arange(N*N).reshape(N,N).contiguous().realize()
|
||||
b = Tensor.arange(N*N).reshape(N,N).contiguous().realize()
|
||||
a = Tensor.arange(N*N).reshape(N,N).clone().realize()
|
||||
b = Tensor.arange(N*N).reshape(N,N).clone().realize()
|
||||
new_a = (a.T+b).numpy()
|
||||
a.assign(a.T+b)
|
||||
np.testing.assert_allclose(a.numpy(), new_a)
|
||||
|
||||
def test_post_flipped_assignment(self):
|
||||
a = Tensor.arange(N*N).reshape(N,N).contiguous().realize()
|
||||
b = Tensor.arange(N*N).reshape(N,N).contiguous().realize()
|
||||
a = Tensor.arange(N*N).reshape(N,N).clone().realize()
|
||||
b = Tensor.arange(N*N).reshape(N,N).clone().realize()
|
||||
new_a = (a.flip(0)+b).numpy()
|
||||
a.assign(a.flip(0)+b)
|
||||
np.testing.assert_allclose(a.numpy(), new_a)
|
||||
|
||||
def test_post_flipped_assignment_axis1(self):
|
||||
a = Tensor.arange(N*N).reshape(N,N).contiguous().realize()
|
||||
b = Tensor.arange(N*N).reshape(N,N).contiguous().realize()
|
||||
a = Tensor.arange(N*N).reshape(N,N).clone().realize()
|
||||
b = Tensor.arange(N*N).reshape(N,N).clone().realize()
|
||||
new_a = (a.flip(1)+b).numpy()
|
||||
a.assign(a.flip(1)+b)
|
||||
np.testing.assert_allclose(a.numpy(), new_a)
|
||||
|
||||
def test_post_reshape_assignment_fine(self):
|
||||
a = Tensor.arange(N*N).reshape(N, N).contiguous().realize()
|
||||
b = Tensor.arange(N*N).reshape(N, N).contiguous().realize()
|
||||
a = Tensor.arange(N*N).reshape(N, N).clone().realize()
|
||||
b = Tensor.arange(N*N).reshape(N, N).clone().realize()
|
||||
rhs = a.reshape(-1).reshape(N, N)
|
||||
new_a = (rhs+b).numpy()
|
||||
a.assign(rhs+b) # self-assign with reshape view is fine
|
||||
@@ -355,7 +355,7 @@ class TestAssign(unittest.TestCase):
|
||||
|
||||
@unittest.skip("multi output not supported anymore")
|
||||
def test_simple_assignment_multioutput(self):
|
||||
a = Tensor.arange(32*32).reshape(32, 32).contiguous().realize()
|
||||
a = Tensor.arange(32*32).reshape(32, 32).clone().realize()
|
||||
b = Tensor.full((32, ), 1.).contiguous().realize()
|
||||
c = Tensor.full((32, ), 2.).contiguous().realize()
|
||||
d = Tensor.full((32, ), 3.).contiguous().realize()
|
||||
@@ -375,15 +375,15 @@ class TestAssign(unittest.TestCase):
|
||||
# NOTE: if the assign target is read/write in a single kernel, it should be contiguous
|
||||
|
||||
def test_permuted_assignment_correct(self):
|
||||
a = Tensor.arange(4 * 4).reshape(4, 4).contiguous().realize()
|
||||
b = Tensor.arange(4 * 4).reshape(4, 4).contiguous().realize()
|
||||
a = Tensor.arange(4 * 4).reshape(4, 4).clone().realize()
|
||||
b = Tensor.arange(4 * 4).reshape(4, 4).clone().realize()
|
||||
a = a.permute(1, 0)
|
||||
new_val = a + b
|
||||
a.assign(new_val)
|
||||
np.testing.assert_equal(a.numpy(), np.arange(4 * 4).reshape(4, 4).transpose(1, 0) + np.arange(4 * 4).reshape(4, 4))
|
||||
|
||||
def test_permuted_reduceop_child_dual_use(self):
|
||||
a = Tensor.arange(32*32*32).reshape(32, 32, 32).contiguous().realize()
|
||||
a = Tensor.arange(32*32*32).reshape(32, 32, 32).clone().realize()
|
||||
b = Tensor.ones(32, 32, dtype=dtypes.int).contiguous().realize()
|
||||
r = a.sum(axis=1)
|
||||
b.assign(r + b.permute(1, 0))
|
||||
@@ -392,7 +392,7 @@ class TestAssign(unittest.TestCase):
|
||||
|
||||
@unittest.skip("multi output not supported anymore")
|
||||
def test_permuted_reduceop_multioutput_dual_use(self):
|
||||
a = Tensor.arange(32*32*32).reshape(32, 32, 32).contiguous().realize()
|
||||
a = Tensor.arange(32*32*32).reshape(32, 32, 32).clone().realize()
|
||||
b = Tensor.full((32, 32), 1.).contiguous().realize()
|
||||
c = Tensor.full((32, 32), 2.).contiguous().realize()
|
||||
|
||||
@@ -405,9 +405,9 @@ class TestAssign(unittest.TestCase):
|
||||
|
||||
@unittest.skip("multi output not supported anymore")
|
||||
def test_permuted_reduceop_multioutput_dual_use_possible(self):
|
||||
a = Tensor.arange(32*32*32).reshape(32, 32, 32).contiguous().realize()
|
||||
b = Tensor.arange(32 * 32).reshape(32, 32).realize()
|
||||
c = Tensor.arange(32 * 32).reshape(32, 32).realize()
|
||||
a = Tensor.arange(32*32*32).reshape(32, 32, 32).clone().realize()
|
||||
b = Tensor.arange(32 * 32).reshape(32, 32).clone().realize()
|
||||
c = Tensor.arange(32 * 32).reshape(32, 32).clone().realize()
|
||||
|
||||
GlobalCounters.reset()
|
||||
r = a.sum(axis=1)
|
||||
@@ -441,7 +441,7 @@ class TestAssign(unittest.TestCase):
|
||||
# Forward shift: read index > write index in overlap
|
||||
N = 100000
|
||||
shift = 1000
|
||||
a = Tensor.arange(N).float().contiguous().realize()
|
||||
a = Tensor.arange(N).float().clone().realize()
|
||||
expected = np.arange(N, dtype=np.float32)
|
||||
expected[:N-shift] = expected[shift:].copy()
|
||||
with Context(NOOPT=1): a[0:N-shift].assign(a[shift:N]).realize()
|
||||
@@ -451,7 +451,7 @@ class TestAssign(unittest.TestCase):
|
||||
# Reverse shift: write index > read index in overlap
|
||||
N = 100000
|
||||
shift = 1000
|
||||
a = Tensor.arange(N).float().contiguous().realize()
|
||||
a = Tensor.arange(N).float().clone().realize()
|
||||
expected = np.arange(N, dtype=np.float32)
|
||||
expected[shift:] = expected[:N-shift].copy()
|
||||
with Context(NOOPT=1): a[shift:N].assign(a[0:N-shift]).realize()
|
||||
@@ -459,7 +459,7 @@ class TestAssign(unittest.TestCase):
|
||||
|
||||
def test_nonoverlapping_shrink_assignment(self):
|
||||
# TODO: non-overlapping shrinks don't actually need contiguous, could be 1 kernel with smarter range analysis
|
||||
a = Tensor.arange(100).float().contiguous().realize()
|
||||
a = Tensor.arange(100).float().clone().realize()
|
||||
expected = np.arange(100, dtype=np.float32)
|
||||
expected[0:10] = expected[50:60].copy()
|
||||
GlobalCounters.reset()
|
||||
|
||||
Reference in New Issue
Block a user