github/tinygrad
1
1
Fork 0
mirror of https://github.com/tinygrad/tinygrad.git synced 2026-04-07 03:00:26 -04:00
Code Issues Packages Projects Releases Wiki Activity

symbolic Tensor.var (#4843)

Browse Source 
taken from #4446 and add more tests
This commit is contained in:
chenyu
2024-06-05 12:55:54 -04:00
committed by GitHub
parent 085c0bbf6b
commit a352b6d9ce
4 changed files with 100 additions and 52 deletions
Download Patch File Download Diff File Expand all files Collapse all files

53
test/test_symbolic_jit.py
View File

@@ -241,5 +241,58 @@ class TestSymbolicJit(unittest.TestCase):
expected = a.mean(1).numpy()
np.testing.assert_allclose(symbolic, expected, atol=1e-6, rtol=1e-6)
def test_var(self):
def f(a): return a.var().realize()
def f0(a): return a.var(0).realize()
def f1(a): return a.var(1).realize()
jf = TinyJit(f)
jf0 = TinyJit(f0)
jf1 = TinyJit(f1)
for i in range(1, 5):
vi = Variable("i", 1, 10).bind(i)
# aixs = None
a = Tensor.rand(i, 3)
symbolic = jf(a.reshape(vi, 3)).numpy()
expected = a.var().numpy()
np.testing.assert_allclose(symbolic, expected, atol=1e-6, rtol=1e-6)
# aixs = 0
a = Tensor.rand(i, 3)
symbolic = jf0(a.reshape(vi, 3)).numpy()
expected = a.var(0).numpy()
np.testing.assert_allclose(symbolic, expected, atol=1e-6, rtol=1e-6)
# aixs = 1
a = Tensor.rand(i, 3)
symbolic = jf1(a.reshape(vi, 3)).reshape(i).numpy()
expected = a.var(1).numpy()
np.testing.assert_allclose(symbolic, expected, atol=1e-6, rtol=1e-6)
@unittest.skip("failed for some")
def test_var_2d(self):
def f(a): return a.var().realize()
def f0(a): return a.var(0).realize()
def f1(a): return a.var(1).realize()
jf = TinyJit(f)
jf0 = TinyJit(f0)
jf1 = TinyJit(f1)
for i in range(1, 5):
for j in range(1, 5):
vi = Variable("i", 1, 10).bind(i)
vj = Variable("j", 1, 10).bind(j)
# aixs = None
a = Tensor.rand(i, j)
symbolic = jf(a.reshape(vi, vj)).numpy()
expected = a.var().numpy()
np.testing.assert_allclose(symbolic, expected, atol=1e-6, rtol=1e-6)
# aixs = 0
a = Tensor.rand(i, j)
symbolic = jf0(a.reshape(vi, vj)).reshape(j).numpy()
expected = a.var(0).numpy()
np.testing.assert_allclose(symbolic, expected, atol=1e-6, rtol=1e-6)
# aixs = 1
a = Tensor.rand(i, j)
symbolic = jf1(a.reshape(vi, vj)).reshape(i).numpy()
expected = a.var(1).numpy()
np.testing.assert_allclose(symbolic, expected, atol=1e-6, rtol=1e-6)
if __name__ == '__main__':
unittest.main()

58
test/test_symbolic_ops.py
View File

@@ -152,42 +152,42 @@ class TestSymbolicOps(unittest.TestCase):
def test_mean(self):
for i in range(1, 5):
vi = Variable("i", 1, 10).bind(i)
# aixs = None
a = Tensor.rand(i, 3)
symbolic = a.reshape(vi, 3).mean().numpy()
expected = a.mean().numpy()
np.testing.assert_allclose(symbolic, expected, atol=1e-6, rtol=1e-6)
# aixs = 0
a = Tensor.rand(i, 3)
symbolic = a.reshape(vi, 3).mean(0).numpy()
expected = a.mean(0).numpy()
np.testing.assert_allclose(symbolic, expected, atol=1e-6, rtol=1e-6)
# aixs = 1
a = Tensor.rand(i, 3)
symbolic = a.reshape(vi, 3).mean(1).reshape(i).numpy()
expected = a.mean(1).numpy()
np.testing.assert_allclose(symbolic, expected, atol=1e-6, rtol=1e-6)
for axis in [None, 0, 1]:
a = Tensor.rand(i, 3)
expected = a.mean(axis).numpy()
symbolic = a.reshape(vi, 3).mean(axis).reshape(expected.shape).numpy()
np.testing.assert_allclose(symbolic, expected, atol=1e-6, rtol=1e-6)
def test_mean_2d(self):
for i in range(1, 5):
for j in range(1, 5):
vi = Variable("i", 1, 10).bind(i)
vj = Variable("j", 1, 10).bind(j)
# aixs = None
a = Tensor.rand(i, j)
symbolic = a.reshape(vi, vj).mean().numpy()
expected = a.mean().numpy()
np.testing.assert_allclose(symbolic, expected, atol=1e-6, rtol=1e-6)
# aixs = 0
a = Tensor.rand(i, j)
symbolic = a.reshape(vi, vj).mean(0).reshape(j).numpy()
expected = a.mean(0).numpy()
np.testing.assert_allclose(symbolic, expected, atol=1e-6, rtol=1e-6)
# aixs = 1
a = Tensor.rand(i, j)
symbolic = a.reshape(vi, vj).mean(1).reshape(i).numpy()
expected = a.mean(1).numpy()
for axis in [None, 0, 1]:
a = Tensor.rand(i, j)
expected = a.mean(axis).numpy()
symbolic = a.reshape(vi, vj).mean(axis).reshape(expected.shape).numpy()
np.testing.assert_allclose(symbolic, expected, atol=1e-6, rtol=1e-6)
def test_var(self):
for i in range(1, 5):
vi = Variable("i", 1, 10).bind(i)
for axis in [None, 0, 1]:
a = Tensor.rand(i, 3)
expected = a.var(axis).numpy()
symbolic = a.reshape(vi, 3).var(axis).reshape(expected.shape).numpy()
np.testing.assert_allclose(symbolic, expected, atol=1e-6, rtol=1e-6)
def test_var_2d(self):
for i in range(1, 5):
for j in range(1, 5):
vi = Variable("i", 1, 10).bind(i)
vj = Variable("j", 1, 10).bind(j)
for axis in [None, 0, 1]:
a = Tensor.rand(i, j)
expected = a.var(axis).numpy()
symbolic = a.reshape(vi, vj).var(axis).reshape(expected.shape).numpy()
np.testing.assert_allclose(symbolic, expected, atol=1e-6, rtol=1e-6)
if __name__ == '__main__':
unittest.main()

29
test/test_tensor_variable.py
View File

@@ -26,38 +26,31 @@ class TestTensorVariable(unittest.TestCase):
assert (Tensor(3) * (vv * 4)).item() == 24
def test_symbolic_mean(self):
vv = Variable("a", 1, 10)
vv.bind(2)
vv = Variable("a", 1, 10).bind(2)
t = Tensor.ones(2, 2).contiguous().reshape(2, vv)
ret = t.mean().item()
assert ret == 1
@unittest.skip("symbolic var isn't supported")
def test_symbolic_var(self):
vv = Variable("a", 1, 10)
vv.bind(2)
t = Tensor.ones(2, 2).contiguous().reshape(2, vv)
ret = t.var().item()
assert ret == 0
def test_symbolic_mean_2d(self):
vv = Variable("a", 1, 10)
vv.bind(2)
vv2 = Variable("b", 1, 10)
vv2.bind(2)
vv = Variable("a", 1, 10).bind(2)
vv2 = Variable("b", 1, 10).bind(2)
t = Tensor.ones(2, 2).contiguous().reshape(vv2, vv)
ret = t.mean().item()
assert ret == 1
def test_symbolic_mean_2d_axis_1(self):
vv = Variable("a", 1, 10)
vv.bind(2)
vv2 = Variable("b", 1, 10)
vv2.bind(2)
vv = Variable("a", 1, 10).bind(2)
vv2 = Variable("b", 1, 10).bind(2)
t = Tensor.ones(2, 2).contiguous().reshape(vv2, vv)
ret = t.mean(axis=1).reshape(2, 1).numpy()
assert np.all(ret == 1)
def test_symbolic_var(self):
vv = Variable("a", 1, 10).bind(2)
t = Tensor.ones(2, 2).contiguous().reshape(2, vv)
ret = t.var().item()
assert ret == 0
@unittest.skip("symbolic arange isn't supported")
def test_symbolic_arange(self):
vv = Variable("a", 1, 10)