openpilot is an open source driver assistance system. openpilot performs the functions of Automated Lane Centering and Adaptive Cruise Control for over 200 supported car makes and models.
You can not select more than 25 topics Topics must start with a letter or number, can include dashes ('-') and can be up to 35 characters long.

66 lines
2.7 KiB

import numpy as np
import unittest
from tinygrad import Tensor
from typing import List
import functools
def orthogonality_helper(A:Tensor,tolerance=1.0e-5):
b_shape,m = A.shape[0:-2],A.shape[-2] #outer dimension should be the dim along orthogonality
A_identity = (Tensor.eye(m).reshape((1,) * len(b_shape)+(m,m)).expand(b_shape+(m,m)))
np.testing.assert_allclose((A @ A.transpose(-2,-1)).numpy(),A_identity.numpy(),atol=tolerance,rtol=tolerance)
def reconstruction_helper(A:List[Tensor],B:Tensor, tolerance=1.0e-5):
reconstructed_tensor = functools.reduce(Tensor.matmul, A)
np.testing.assert_allclose(reconstructed_tensor.numpy(),B.numpy(),atol=tolerance,rtol=tolerance)
class TestLinAlg(unittest.TestCase):
def test_svd_general(self):
sizes = [(2,2),(5,3),(3,5),(2,2,2,2,3)]
for size in sizes:
a = Tensor.randn(size).realize()
U,S,V = Tensor.svd(a)
b_shape,m,n = size[0:-2],size[-2],size[-1]
k = min(m,n)
s_diag = (S.unsqueeze(-2) * Tensor.eye(k).reshape((1,) * len(b_shape) + (k,k)))
s_diag = s_diag.expand(b_shape + (k,k)).pad(tuple([(0,0) for _ in range(len(size)-2)] + [(0,m-k), (0,n-k)]))
orthogonality_helper(U)
orthogonality_helper(V)
reconstruction_helper([U,s_diag,V],a)
def test_svd_nonfull(self):
sizes = [(2,2),(5,3),(3,5),(2,2,2,2,3)]
for size in sizes:
a = Tensor.randn(size).realize()
U,S,V = Tensor.svd(a,full_matrices=False)
b_shape,m,n = size[0:-2],size[-2],size[-1]
k = min(m,n)
s_diag = (S.unsqueeze(-2) * Tensor.eye(k).reshape((1,) * len(b_shape) + (k,k)).expand(b_shape + (k,k)))
#reduced U,V is only orthogonal along smaller dim
if (m < n): orthogonality_helper(U),orthogonality_helper(V)
else: orthogonality_helper(U.transpose(-2,-1)),orthogonality_helper(V.transpose(-2,-1))
reconstruction_helper([U,s_diag,V],a)
@unittest.skip("very big. recommend wrapping with TinyJit around inner function")
def test_svd_large(self):
size = (1024,1024)
a = Tensor.randn(size).realize()
U,S,V = Tensor.svd(a)
b_shape,m,n = size[0:-2],size[-2],size[-1]
k = min(m,n)
s_diag = (S.unsqueeze(-2) * Tensor.eye(k).reshape((1,) * len(b_shape) + (k,k)))
s_diag = s_diag.expand(b_shape + (k,k)).pad(tuple([(0,0) for _ in range(len(size)-2)] + [(0,m-k), (0,n-k)]))
orthogonality_helper(U,tolerance=1.0e-3)
orthogonality_helper(V,tolerance=1.0e-3)
reconstruction_helper([U,s_diag,V],a,tolerance=1.0e-3)
def test_qr_general(self):
sizes = [(3,3),(3,6),(6,3),(2,2,2,2,2)]
for size in sizes:
a = Tensor.randn(size).realize()
Q,R = Tensor.qr(a)
orthogonality_helper(Q)
reconstruction_helper([Q,R],a)
if __name__ == "__main__":
unittest.main()