mirror of
https://github.com/firestar5683/StarPilot.git
synced 2026-10-28 00:53:45 +08:00
Daniel Koepping
215 lines
9.4 KiB
Python
215 lines
9.4 KiB
Python
import unittest, math
|
|
import numpy as np
|
|
from tinygrad import Tensor
|
|
from tinygrad.dtype import dtypes
|
|
from tinygrad.uop.ops import UOp, KernelInfo, Ops
|
|
|
|
class TestTensorGradient(unittest.TestCase):
|
|
def test_example(self):
|
|
x = Tensor.eye(3)
|
|
y = Tensor([[2.0,0,-2.0]])
|
|
z = y.matmul(x).sum()
|
|
dx, dy = z.gradient(x, y)
|
|
self.assertListEqual(dx.tolist(), [[2.0, 2.0, 2.0], [0.0, 0.0, 0.0], [-2.0, -2.0, -2.0]])
|
|
self.assertListEqual(dy.tolist(), [[1.0, 1.0, 1.0]])
|
|
|
|
def test_zero_if_not_used(self):
|
|
x = Tensor([1.0, 2.0, 3.0])
|
|
w = Tensor.randn((3,))
|
|
self.assertListEqual(x.sum().gradient(w)[0].tolist(), [0.0, 0.0, 0.0])
|
|
|
|
def test_with_custom_gradient(self):
|
|
x = Tensor([1.0, 2.0, 3.0])
|
|
z = (x * x).sum()
|
|
dx = z.gradient(x, gradient=Tensor([3.0]))[0]
|
|
self.assertListEqual(dx.tolist(), [6.0, 12.0, 18.0])
|
|
|
|
def test_broadcast_gradient(self):
|
|
x = Tensor([[1.0], [2.0], [3.0]])
|
|
y = Tensor([[10.0, 20.0, 30.0, 40.0]])
|
|
z = (x + y).sum()
|
|
dx, dy = z.gradient(x, y)
|
|
self.assertListEqual(dx.tolist(), [[4.0], [4.0], [4.0]])
|
|
self.assertListEqual(dy.tolist(), [[3.0, 3.0, 3.0, 3.0]])
|
|
|
|
def test_non_scalar_output(self):
|
|
x = Tensor([1.0, 2.0, 3.0])
|
|
z = x * x
|
|
with self.assertRaises(AssertionError): z.gradient(x)
|
|
dz = Tensor([1.0, 1.0, 1.0])
|
|
dx = z.gradient(x, gradient=dz)[0]
|
|
self.assertListEqual(dx.tolist(), [2.0, 4.0, 6.0])
|
|
|
|
def test_copy_to_device_gradient(self):
|
|
t = Tensor([1.0, 2, 3]).realize()
|
|
t.to("CPU:1").square().sum().backward()
|
|
self.assertEqual(t.grad.device, t.device)
|
|
self.assertListEqual(t.grad.tolist(), [2.0, 4.0, 6.0])
|
|
|
|
def test_multi_device_copy_gradient(self):
|
|
for src in ("CPU", ("CPU", "CPU:1")):
|
|
for dst in ("CPU:2", ("CPU:2", "CPU:3")):
|
|
with self.subTest(src=src, dst=dst):
|
|
t = Tensor([1.0, 2, 3]).to(src).realize()
|
|
grad = t.to(dst).square().sum().gradient(t)[0]
|
|
self.assertEqual(grad.device, t.device)
|
|
self.assertListEqual(grad.tolist(), [2.0, 4.0, 6.0])
|
|
|
|
def test_multiple_backward(self):
|
|
x = Tensor([3.])
|
|
(x*2)[0].backward()
|
|
np.testing.assert_allclose(x.grad.numpy(), [2.0])
|
|
old_grad = x.grad
|
|
(x*3)[0].backward()
|
|
np.testing.assert_allclose(x.grad.numpy(), [2.0+3.0])
|
|
self.assertIs(x.grad, old_grad)
|
|
(x*x)[0].backward()
|
|
np.testing.assert_allclose(x.grad.numpy(), [2.0+3.0+2*3.0])
|
|
self.assertIs(x.grad, old_grad)
|
|
|
|
def test_gradient_through_clone_from_grad_src(self):
|
|
# unlike torch, tinygrad accumulates grad on every tensor in the graph, including non-leaf x
|
|
src = Tensor([1.0, 2.0, 3.0, 4.0])
|
|
x = src.clone()
|
|
(x * 2.0).sum().backward()
|
|
np.testing.assert_allclose(src.grad.numpy(), [2.0, 2.0, 2.0, 2.0])
|
|
np.testing.assert_allclose(x.grad.numpy(), [2.0, 2.0, 2.0, 2.0])
|
|
|
|
def test_gradient_through_clone_from_detached_src(self):
|
|
base = Tensor([1.0, 2.0, 3.0, 4.0])
|
|
x = base.detach().clone()
|
|
(x * 2.0).sum().backward()
|
|
np.testing.assert_allclose(x.grad.numpy(), [2.0, 2.0, 2.0, 2.0]) # gradient flows through clone
|
|
np.testing.assert_allclose(base.grad.numpy(), [0.0, 0.0, 0.0, 0.0]) # ...but detach blocks it from base
|
|
|
|
def test_gradient_through_chained_unrealized_setitem(self):
|
|
g1 = Tensor.zeros(4).contiguous()
|
|
g1[2] = Tensor(1.0)
|
|
g2 = Tensor.zeros(5, 4).contiguous()
|
|
g2[0] = g1
|
|
x = Tensor.randn(4, 4)
|
|
np.testing.assert_allclose(x.pad(((1,0),(0,0))).gradient(x, gradient=g2)[0].numpy(), np.zeros((4, 4)))
|
|
|
|
def test_implicit_broadcast_where_gradient(self):
|
|
# WHERE with a bare ()-shape branch: the scalar's gradient counts the positions where it is selected
|
|
cond, x, w = Tensor([True, False, True]), Tensor([1.0, 2.0, 3.0]), Tensor(4.0, dtype=dtypes.float32)
|
|
dw = Tensor(cond.uop.alu(Ops.WHERE, x.uop, w.uop)).sum().gradient(w)[0]
|
|
self.assertEqual(dw.shape, ())
|
|
self.assertEqual(dw.item(), 1.0)
|
|
dw = Tensor(cond.uop.alu(Ops.WHERE, w.uop, x.uop)).sum().gradient(w)[0]
|
|
self.assertEqual(dw.item(), 2.0)
|
|
|
|
def test_implicit_broadcast_alu_gradient(self):
|
|
# MUL with a bare ()-shape src, no EXPAND in the graph
|
|
x, w = Tensor([1.0, 2.0, 3.0]), Tensor(2.0, dtype=dtypes.float32)
|
|
m = x.uop.alu(Ops.MUL, w.uop)
|
|
self.assertIs(m.src[1], w.uop)
|
|
dw = Tensor(m).sum().gradient(w)[0]
|
|
self.assertEqual(dw.shape, ())
|
|
self.assertEqual(dw.item(), 6.0)
|
|
|
|
def test_implicit_broadcast_intermediate_accumulation(self):
|
|
# s is used directly and through an implicit broadcast edge, each edge's gradient reduces to s's shape before they sum
|
|
x, p = Tensor([1.0, 2.0, 3.0]), Tensor(0.5, dtype=dtypes.float32)
|
|
s = p.sin()
|
|
z = Tensor(x.uop.alu(Ops.MUL, s.uop)).sum() + s
|
|
dp = z.gradient(p)[0]
|
|
self.assertEqual(dp.shape, ())
|
|
self.assertAlmostEqual(dp.item(), 7*math.cos(0.5), places=5)
|
|
|
|
def test_max_backward_many_ties(self):
|
|
t = Tensor.ones(70000, dtype=dtypes.half).contiguous()
|
|
np.testing.assert_allclose(t.max().gradient(t)[0].sum().numpy(), 1.0, atol=1e-3)
|
|
|
|
class TestMultiOutputGradient(unittest.TestCase):
|
|
@staticmethod
|
|
def addmul_kernel(C:UOp, D:UOp, A:UOp, B:UOp) -> UOp:
|
|
C, D, A, B = C.flatten(), D.flatten(), A.flatten(), B.flatten()
|
|
i = UOp.range(C.numel(), 0)
|
|
store_c = C[i].store(A[i] + B[i])
|
|
store_d = D[i].store(A[i] * B[i])
|
|
return UOp.group(store_c, store_d).end(i).sink(arg=KernelInfo(name="addmul")).simplify()
|
|
@staticmethod
|
|
def backward_addmul(grad_c, grad_d, call):
|
|
_c, _d, a, b = call.src[1:]
|
|
grad_a = (Tensor(grad_c) + Tensor(grad_d) * Tensor(b)).uop
|
|
grad_b = (Tensor(grad_c) + Tensor(grad_d) * Tensor(a)).uop
|
|
return (None, None, grad_a, grad_b)
|
|
|
|
def test_custom_kernel_multi_output_backward(self):
|
|
a_np, b_np = np.random.randn(4, 4).astype(np.float32), np.random.randn(4, 4).astype(np.float32)
|
|
a_ref, b_ref = Tensor(a_np), Tensor(b_np)
|
|
((a_ref + b_ref).sum() + (a_ref * b_ref).sum()).backward()
|
|
|
|
a, b = Tensor(a_np), Tensor(b_np)
|
|
Tensor.realize(a, b)
|
|
c, d, _, _ = Tensor.custom_kernel(Tensor.empty(4, 4), Tensor.empty(4, 4), a, b, fxn=self.addmul_kernel, grad_fxn=self.backward_addmul)
|
|
(c.sum() + d.sum()).backward()
|
|
np.testing.assert_allclose(a.grad.numpy(), a_ref.grad.numpy(), rtol=1e-5)
|
|
np.testing.assert_allclose(b.grad.numpy(), b_ref.grad.numpy(), rtol=1e-5)
|
|
|
|
def test_custom_kernel_aliased_output_views_backward(self):
|
|
def kernel(c:UOp, d:UOp, a:UOp) -> UOp:
|
|
c, d, a = c.flatten(), d.flatten(), a.flatten()
|
|
i = UOp.range(2, 0)
|
|
return UOp.group(c[i].store(a[i] * 2), d[i].store(a[i] * 3)).end(i).sink(arg=KernelInfo(name="aliased_outputs"))
|
|
def backward(grad_c:UOp, call:UOp): return (None, None, grad_c)
|
|
|
|
a = Tensor([1., 2.]).contiguous().realize()
|
|
a.requires_grad = True
|
|
out = Tensor.empty(4).contiguous().realize()
|
|
c, _, _ = Tensor.custom_kernel(out[:2], out[2:], a, fxn=kernel, grad_fxn=backward)
|
|
c.sum().backward()
|
|
np.testing.assert_equal(a.grad.numpy(), [1., 1.])
|
|
|
|
def test_custom_kernel_multi_output_backward_interacting(self):
|
|
a_np, b_np = np.random.randn(4, 4).astype(np.float32), np.random.randn(4, 4).astype(np.float32)
|
|
a_ref, b_ref = Tensor(a_np), Tensor(b_np)
|
|
((a_ref + b_ref) * (a_ref * b_ref)).sum().backward()
|
|
|
|
a, b = Tensor(a_np), Tensor(b_np)
|
|
Tensor.realize(a, b)
|
|
c, d, _, _ = Tensor.custom_kernel(Tensor.empty(4, 4), Tensor.empty(4, 4), a, b, fxn=self.addmul_kernel, grad_fxn=self.backward_addmul)
|
|
(c * d).sum().backward()
|
|
np.testing.assert_allclose(a.grad.numpy(), a_ref.grad.numpy(), rtol=1e-5, atol=1e-7)
|
|
np.testing.assert_allclose(b.grad.numpy(), b_ref.grad.numpy(), rtol=1e-5, atol=1e-7)
|
|
|
|
def test_custom_kernel_three_output_backward(self):
|
|
def addmulsub_kernel(C:UOp, D:UOp, E:UOp, A:UOp, B:UOp) -> UOp:
|
|
C, D, E, A, B = C.flatten(), D.flatten(), E.flatten(), A.flatten(), B.flatten()
|
|
i = UOp.range(C.numel(), 0)
|
|
store_c = C[i].store(A[i] + B[i])
|
|
store_d = D[i].store(A[i] * B[i])
|
|
store_e = E[i].store(A[i] - B[i])
|
|
return UOp.group(store_c, store_d, store_e).end(i).sink(arg=KernelInfo(name="addmulsub")).simplify()
|
|
def backward_addmulsub(grad_c, grad_d, grad_e, call):
|
|
_c, _d, _e, a, b = call.src[1:]
|
|
grad_a = (Tensor(grad_c) + Tensor(grad_d) * Tensor(b) + Tensor(grad_e)).uop
|
|
grad_b = (Tensor(grad_c) + Tensor(grad_d) * Tensor(a) - Tensor(grad_e)).uop
|
|
return (None, None, None, grad_a, grad_b)
|
|
|
|
a_np, b_np = np.random.randn(4, 4).astype(np.float32), np.random.randn(4, 4).astype(np.float32)
|
|
a_ref, b_ref = Tensor(a_np), Tensor(b_np)
|
|
((a_ref + b_ref).sum() + (a_ref * b_ref).sum() + (a_ref - b_ref).sum()).backward()
|
|
|
|
a, b = Tensor(a_np), Tensor(b_np)
|
|
Tensor.realize(a, b)
|
|
c, d, e, _, _ = Tensor.custom_kernel(Tensor.empty(4, 4), Tensor.empty(4, 4), Tensor.empty(4, 4), a, b,
|
|
fxn=addmulsub_kernel, grad_fxn=backward_addmulsub)
|
|
(c.sum() + d.sum() + e.sum()).backward()
|
|
np.testing.assert_allclose(a.grad.numpy(), a_ref.grad.numpy(), atol=1e-6, rtol=1e-5)
|
|
np.testing.assert_allclose(b.grad.numpy(), b_ref.grad.numpy(), atol=1e-6, rtol=1e-5)
|
|
|
|
class TestViewGradient(unittest.TestCase):
|
|
def test_expand(self):
|
|
x = Tensor.randn(5,2)
|
|
a = Tensor([3.])
|
|
aex = a.expand(10)
|
|
(aex.reshape(5,2) * x).sum().backward()
|
|
np.testing.assert_allclose(aex.grad.numpy(), x.reshape(10).numpy())
|
|
with self.assertRaises(AssertionError):
|
|
np.testing.assert_allclose(aex.grad.numpy(), a.grad.expand(10).numpy())
|
|
|
|
if __name__ == '__main__':
|
|
unittest.main()
|