mirror of
https://github.com/firestar5683/StarPilot.git
synced 2026-10-28 00:53:45 +08:00
Daniel Koepping
95 lines
3.3 KiB
Python
95 lines
3.3 KiB
Python
import unittest
|
|
from tinygrad import Tensor
|
|
|
|
class TestAfterCounterexamples(unittest.TestCase):
|
|
def test_ordered_writes_allowed(self):
|
|
x = Tensor([0.]).realize().uop
|
|
a = x.after(x.store(1))
|
|
b = a.after(a.store(2))
|
|
self.assertEqual(Tensor(b).tolist(), [2.])
|
|
|
|
def test_disjoint_writes_allowed(self):
|
|
x = Tensor([0., 0.]).realize().uop
|
|
y = Tensor(x.after(x[:1].store(1), x[1:].store(2)))
|
|
self.assertEqual(y.tolist(), [1., 2.])
|
|
|
|
def test_read_modify_write_chain(self):
|
|
x = Tensor([2.]).clone()
|
|
x.assign(x + 1)
|
|
x.assign(x * 2)
|
|
self.assertEqual(x.tolist(), [6.])
|
|
|
|
def test_overwrite_cuts_gradient(self):
|
|
x = Tensor([2.])
|
|
y = x.clone()
|
|
y.assign(3) # overwriting with a constant makes y independent of x
|
|
self.assertEqual(y.sum().gradient(x)[0].tolist(), [0.])
|
|
|
|
def test_shared_state_readers(self):
|
|
x = Tensor([2.]).clone()
|
|
x.assign(x + 1)
|
|
a, b = x + 1, x * 2
|
|
Tensor.realize(a, b)
|
|
self.assertEqual(a.tolist(), [4.])
|
|
self.assertEqual(b.tolist(), [6.])
|
|
|
|
@unittest.expectedFailure
|
|
def test_chained_square_assign_gradient(self):
|
|
x = Tensor([2.0])
|
|
y = x.clone()
|
|
y.assign(y*y)
|
|
y.assign(y*y)
|
|
# y = x**4, so dy/dx = 4*x**3. Currently raises "cycle detected while indexing".
|
|
self.assertEqual(y.sum().gradient(x)[0].tolist(), [32.])
|
|
|
|
def test_partial_store_gradient(self):
|
|
x = Tensor([2., 3.]).realize()
|
|
y = Tensor(x.uop.after(x[:1].uop.store(4)))
|
|
# y = [4, x[1]]; only the untouched element depends on x.
|
|
self.assertEqual(y.sum().gradient(x)[0].tolist(), [0., 1.])
|
|
|
|
def test_partial_store_source_gradient(self):
|
|
x = Tensor([4.])
|
|
y = Tensor([2., 3.]).realize()
|
|
z = Tensor(y.uop.after(y[:1].uop.store(x.uop)))
|
|
# x contributes once, not twice.
|
|
self.assertEqual(z.sum().gradient(x)[0].tolist(), [1.])
|
|
|
|
def test_unrelated_store_gradient(self):
|
|
x = Tensor([2.]).realize()
|
|
y = x.clone()
|
|
z = Tensor(x.uop.after(y.uop.store(0)))
|
|
# Zeroing y does not change x.
|
|
self.assertEqual(z.sum().gradient(x)[0].tolist(), [1.])
|
|
|
|
def test_after_dependency_gradient(self):
|
|
x = Tensor([2., 3.])
|
|
y = x.clone()
|
|
y[:1].assign(0)
|
|
# View assign is an AFTER on a partial STORE; only the untouched element depends on x.
|
|
self.assertEqual(y.sum().gradient(x)[0].tolist(), [0., 1.])
|
|
|
|
def test_view_assign_gradient(self):
|
|
for view, expected in ((lambda t: t.reshape(3, 2)[1:], [[1., 1., 0.], [0., 0., 0.]]),
|
|
(lambda t: t.permute(1, 0)[1:], [[1., 0., 0.], [1., 0., 0.]]),
|
|
(lambda t: t.flip((0, 1))[:1], [[1., 1., 1.], [0., 0., 0.]]),
|
|
(lambda t: t[:, ::2], [[0., 1., 0.], [0., 1., 0.]])):
|
|
with self.subTest(expected=expected):
|
|
x = Tensor([[1., 2., 3.], [4., 5., 6.]])
|
|
y = x.clone()
|
|
v = Tensor.full(view(y).shape, 7.)
|
|
view(y).assign(v)
|
|
gx, gv = y.sum().gradient(x, v)
|
|
self.assertEqual(gx.tolist(), expected)
|
|
self.assertEqual(gv.tolist(), Tensor.ones(v.shape).tolist())
|
|
|
|
@unittest.expectedFailure
|
|
def test_unordered_overlapping_stores_rejected(self):
|
|
x = Tensor([0.]).realize().uop
|
|
# No ordering between the writes. Currently succeeds with [2.].
|
|
with self.assertRaises(RuntimeError):
|
|
Tensor(x.after(x.store(1), x.store(2))).realize()
|
|
|
|
if __name__ == "__main__":
|
|
unittest.main()
|