mirror of
https://github.com/firestar5683/StarPilot.git
synced 2026-11-08 21:43:44 +08:00
firestar5683
22a8575199
Keep official openpilot ancestry and dependency licenses. Store resources directly in Git, remove LFS setup requirements, and leave the downloadable big model outside the source tree.
233 lines
9.8 KiB
Python
233 lines
9.8 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_cast_before_view(self):
|
|
x = Tensor([1.0, 1, 1, 1])
|
|
x_reshaped = x.reshape(2,2)
|
|
x_casted = x_reshaped.cast(dtypes.float16)
|
|
x_casted.mean().gradient(x_reshaped)
|
|
|
|
def test_non_float_tensor_raise(self):
|
|
x = Tensor([1, 2, 3])
|
|
with self.assertRaises(RuntimeError): x.sum().gradient(x)
|
|
with self.assertRaises(RuntimeError): x.float().sum().gradient(x)
|
|
|
|
def test_const_target_raise(self):
|
|
t = Tensor(2.0)
|
|
with self.assertRaises(RuntimeError): (t * 2.0).gradient(t)
|
|
|
|
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_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_setitem_on_grad_used_tensor_raises(self):
|
|
x = Tensor([1.0, 2.0, 3.0, 4.0]).realize()
|
|
_ = (x * 2.0).sum()
|
|
with self.assertRaises(RuntimeError):
|
|
x[0] = 99.0
|
|
|
|
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_bare_const_skipped_by_backward(self):
|
|
Tensor.manual_seed(0)
|
|
w = Tensor(1.0)
|
|
(Tensor.rand(()) + w).backward()
|
|
self.assertIsNone(w.grad)
|
|
|
|
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()
|