mirror of
https://github.com/firestar5683/StarPilot.git
synced 2026-10-28 09:03:46 +08:00
Daniel Koepping
79 lines
3.0 KiB
Python
79 lines
3.0 KiB
Python
import unittest
|
|
from tinygrad.function import function
|
|
from tinygrad import Tensor, GlobalCounters
|
|
from tinygrad.uop.ops import UOp, KernelInfo
|
|
|
|
class TestFunction(unittest.TestCase):
|
|
def test_depth_restored_on_exception(self):
|
|
from tinygrad.function import _function
|
|
@function
|
|
def f(a:Tensor) -> Tensor: raise ValueError("error")
|
|
with self.assertRaises(ValueError): f(Tensor([1]))
|
|
self.assertEqual(_function.depth, 0)
|
|
|
|
def test_name(self):
|
|
@function
|
|
def f(a:Tensor) -> Tensor: return a + 1
|
|
assert f(Tensor([1])).uop.src[1].arg.name.endswith("f")
|
|
|
|
def test_method_name(self):
|
|
class Foo:
|
|
@function
|
|
def __call__(self, x:Tensor) -> Tensor: return x + 1
|
|
assert Foo()(Tensor([1])).uop.src[1].arg.name.endswith("Foo.__call__")
|
|
|
|
class TestFunctionMulti(unittest.TestCase):
|
|
devices_2 = ("NULL:0", "NULL:1")
|
|
|
|
def test_call_axis(self):
|
|
@function
|
|
def f(x:Tensor, w:Tensor) -> Tensor: return x @ w
|
|
|
|
x = Tensor([[1.,0.],[0.,1.],[1.,1.],[0.,0.]]).shard(self.devices_2, axis=0)
|
|
w = Tensor([[1.,2.],[3.,4.]]).shard(self.devices_2, axis=None)
|
|
result = f(x, w)
|
|
# CALL output should inherit axis=0 from the sharded input
|
|
self.assertEqual(result.uop.axis, 0)
|
|
# reduce on the sharded axis should remove it
|
|
self.assertIsNone(result.sum().uop.axis)
|
|
|
|
class TestFunctionTuple(unittest.TestCase):
|
|
def test_custom_kernel_inplace_output_is_implicit(self):
|
|
# caller-owned storage must be captured, even before its Buffer is bound
|
|
state = Tensor.empty(4)
|
|
def inplace_add(C:UOp, A:UOp) -> UOp:
|
|
i = UOp.range(A.shape[0], 0)
|
|
return C[i].store(C[i].load() + A[i]).end(i).sink(arg=KernelInfo(name="inplace_add"))
|
|
@function(precompile=True, allow_implicit=False)
|
|
def f(a:Tensor): return Tensor.custom_kernel(state, a, fxn=inplace_add)[0]
|
|
with self.assertRaisesRegex(RuntimeError, "implicit buffer"): f(Tensor([1., 2., 3., 4.]).contiguous().realize())
|
|
|
|
class TestFunctionGrad(unittest.TestCase):
|
|
def test_function_grad_ops(self, precompile=False, precompile_backward=False):
|
|
N = 64
|
|
x = Tensor.ones(N,N).contiguous()
|
|
w1 = Tensor.ones(N,N).contiguous()
|
|
w2 = Tensor.ones(N,N).contiguous()
|
|
w3 = Tensor.ones(N,N).contiguous()
|
|
ref = Tensor.ones(N,N).contiguous()
|
|
Tensor.realize(x, w1, w2, w3, ref)
|
|
@function(precompile=precompile, precompile_backward=precompile_backward)
|
|
def f(x, w1, w2, w3) -> tuple[Tensor, ...]:
|
|
p1 = x@w1
|
|
p2 = p1@w2
|
|
p3 = p2@w3
|
|
return p1, p2, p3, p3.contiguous()
|
|
ret = f(x, w1, w2, w3)[-1]
|
|
loss = (ret-ref).square().mean().backward()
|
|
print("RESET")
|
|
GlobalCounters.reset()
|
|
loss.realize(w1.grad, w2.grad, w3.grad)
|
|
print(GlobalCounters.global_ops, GlobalCounters.global_mem)
|
|
self.assertLessEqual(GlobalCounters.global_ops, 5000000)
|
|
def test_function_grad_ops_precompile(self): self.test_function_grad_ops(precompile=True)
|
|
def test_function_grad_ops_precompile_backward(self):
|
|
self.test_function_grad_ops(precompile=True, precompile_backward=True)
|
|
|
|
if __name__ == '__main__':
|
|
unittest.main()
|