mirror of https://github.com/Jittor/Jittor
261 lines
8.3 KiB
Python
261 lines
8.3 KiB
Python
# ***************************************************************
|
|
# Copyright (c) 2020 Jittor. All Rights Reserved.
|
|
# Authors:
|
|
# Haoyang Peng <2247838039@qq.com>
|
|
# Guowei Yang <471184555@qq.com>
|
|
# Dun Liang <randonlang@gmail.com>.
|
|
#
|
|
# This file is subject to the terms and conditions defined in
|
|
# file 'LICENSE.txt', which is part of this source code package.
|
|
# ***************************************************************
|
|
import jittor as jt
|
|
import numpy as np
|
|
import unittest
|
|
|
|
|
|
try:
|
|
import autograd.numpy as anp
|
|
from autograd import jacobian
|
|
has_autograd = True
|
|
except:
|
|
has_autograd = False
|
|
|
|
@unittest.skipIf(not has_autograd, "No autograd found.")
|
|
class TestCodeOp(unittest.TestCase):
|
|
def test_svd(self):
|
|
def check_svd(a):
|
|
u,s,v = anp.linalg.svd(a, full_matrices=0)
|
|
return u,s,v
|
|
|
|
def check_u(a):
|
|
u,s,v = anp.linalg.svd(a, full_matrices=0)
|
|
return u
|
|
|
|
def check_s(a):
|
|
u,s,v = anp.linalg.svd(a, full_matrices=0)
|
|
return s
|
|
|
|
def check_v(a):
|
|
u,s,v = anp.linalg.svd(a, full_matrices=0)
|
|
return v
|
|
|
|
for i in range(50):
|
|
#not for full-matrices!
|
|
a = jt.random((2,2,5,4))
|
|
c_a = anp.array(a.data)
|
|
u,s,v = jt.linalg.svd(a)
|
|
tu,ts,tv = check_svd(c_a)
|
|
assert np.allclose(tu,u.data)
|
|
assert np.allclose(ts,s.data)
|
|
assert np.allclose(tv,v.data)
|
|
ju = jt.grad(u,a)
|
|
js = jt.grad(s,a)
|
|
jv = jt.grad(v,a)
|
|
grad_u = jacobian(check_u)
|
|
gu = grad_u(c_a)
|
|
gu = np.sum(gu, 4)
|
|
gu = np.sum(gu, 4)
|
|
gu = np.sum(gu, 2)
|
|
gu = np.sum(gu, 2)
|
|
grad_s = jacobian(check_s)
|
|
gs = grad_s(c_a)
|
|
gs = np.sum(gs, 4)
|
|
gs = np.sum(gs, 2)
|
|
gs = np.sum(gs, 2)
|
|
grad_v = jacobian(check_v)
|
|
gv = grad_v(c_a)
|
|
gv = np.sum(gv, 4)
|
|
gv = np.sum(gv, 4)
|
|
gv = np.sum(gv, 2)
|
|
gv = np.sum(gv, 2)
|
|
try:
|
|
assert np.allclose(ju.data,gu,atol=1e-5)
|
|
except AssertionError:
|
|
print(ju.data)
|
|
print(gu)
|
|
try:
|
|
assert np.allclose(js.data,gs,atol=1e-5)
|
|
except AssertionError:
|
|
print(js.data)
|
|
print(gs)
|
|
try:
|
|
assert np.allclose(jv.data,gv,atol=1e-5)
|
|
except AssertionError:
|
|
print(jv.data)
|
|
print(gv)
|
|
|
|
def test_eigh(self):
|
|
def check_eigh(a,UPLO='L'):
|
|
w, v = anp.linalg.eigh(a,UPLO)
|
|
return w, v
|
|
|
|
def check_w(a,UPLO='L'):
|
|
w, v = anp.linalg.eigh(a,UPLO)
|
|
return w
|
|
|
|
def check_v(a,UPLO='L'):
|
|
w, v = anp.linalg.eigh(a,UPLO)
|
|
return v
|
|
|
|
for i in range(50):
|
|
a = jt.random((2,2,3,3))
|
|
c_a = a.data
|
|
w, v = jt.linalg.eigh(a)
|
|
tw, tv = check_eigh(c_a)
|
|
assert np.allclose(w.data,tw)
|
|
assert np.allclose(v.data,tv)
|
|
jw = jt.grad(w, a)
|
|
jv = jt.grad(v, a)
|
|
check_gw = jacobian(check_w)
|
|
check_gv = jacobian(check_v)
|
|
gw = check_gw(c_a)
|
|
gw = np.sum(gw,4)
|
|
gw = np.sum(gw,2)
|
|
gw = np.sum(gw,2)
|
|
assert np.allclose(gw,jw.data,rtol = 1,atol = 5e-8)
|
|
gv = check_gv(c_a)
|
|
gv = np.sum(gv,4)
|
|
gv = np.sum(gv,4)
|
|
gv = np.sum(gv,2)
|
|
gv = np.sum(gv,2)
|
|
assert np.allclose(gv,jv.data,rtol = 1,atol = 5e-8)
|
|
|
|
def test_pinv(self):
|
|
def check_pinv(a):
|
|
w = anp.linalg.pinv(a)
|
|
return w
|
|
|
|
for i in range(50):
|
|
x = jt.random((2,2,4,4))
|
|
c_a = x.data
|
|
mx = jt.linalg.pinv(x)
|
|
tx = check_pinv(c_a)
|
|
np.allclose(mx.data,tx)
|
|
jx = jt.grad(mx,x)
|
|
check_grad = jacobian(check_pinv)
|
|
gx = check_grad(c_a)
|
|
np.allclose(gx,jx.data)
|
|
|
|
def test_inv(self):
|
|
def check_inv(a):
|
|
w = anp.linalg.inv(a)
|
|
return w
|
|
for i in range(50):
|
|
tn = np.random.randn(4,4).astype('float32')*5
|
|
while np.allclose(np.linalg.det(tn),0):
|
|
tn = np.random.randn((4,4)).astype('float32')*5
|
|
x = jt.array(tn)
|
|
x = x.reindex([2,2,x.shape[0],x.shape[1]],["i2","i3"])
|
|
c_a = x.data
|
|
mx = jt.linalg.inv(x)
|
|
tx = check_inv(c_a)
|
|
np.allclose(mx.data,tx)
|
|
jx = jt.grad(mx,x)
|
|
check_grad = jacobian(check_inv)
|
|
gx = check_grad(c_a)
|
|
np.allclose(gx,jx.data)
|
|
|
|
def test_slogdet(self):
|
|
def check_ans(a):
|
|
s, w = anp.linalg.slogdet(a)
|
|
return s, w
|
|
|
|
def check_slogdet(a):
|
|
s, w = anp.linalg.slogdet(a)
|
|
return w
|
|
|
|
for i in range(50):
|
|
tn = np.random.randn(4,4).astype('float32')*10
|
|
while np.allclose(np.linalg.det(tn),0):
|
|
tn = np.random.randn((4,4)).astype('float32')*10
|
|
x = jt.array(tn)
|
|
x = x.reindex([2,2,x.shape[0],x.shape[1]],["i2","i3"])
|
|
s = list(x.shape)
|
|
det_s = s[:-2]
|
|
if len(det_s) == 0:
|
|
det_s.append(1)
|
|
sign, mx = jt.linalg.slogdet(x)
|
|
ts, ta = check_ans(x.data)
|
|
assert np.allclose(sign.data, ts)
|
|
assert np.allclose(mx.data, ta)
|
|
jx = jt.grad(mx,x)
|
|
check_sgrad = jacobian(check_slogdet)
|
|
gx = check_sgrad(x.data)
|
|
gx = np.sum(gx,2)
|
|
gx = np.sum(gx,2)
|
|
assert np.allclose(gx,jx.data)
|
|
|
|
def test_cholesky(self):
|
|
def check_cholesky(a):
|
|
L = anp.linalg.cholesky(a)
|
|
return L
|
|
|
|
for i in range(50):
|
|
x = jt.array(np.diag((np.random.rand(3) + 1) * 2))
|
|
x = x.reindex([2,2,x.shape[0],x.shape[1]],["i2","i3"])
|
|
tx = x.data
|
|
L = jt.linalg.cholesky(x)
|
|
tL = check_cholesky(tx)
|
|
assert np.allclose(tL,L.data)
|
|
jx = jt.grad(L,x)
|
|
check_grad = jacobian(check_cholesky)
|
|
gx = check_grad(tx)
|
|
gx = np.sum(gx, 0)
|
|
gx = np.sum(gx, 0)
|
|
gx = np.sum(gx, 0)
|
|
gx = np.sum(gx, 0)
|
|
assert np.allclose(jx.data,gx)
|
|
|
|
def test_solve(self):
|
|
def check_solve(a,b):
|
|
ans = anp.linalg.solve(a,b)
|
|
return ans
|
|
|
|
for i in range(50):
|
|
a = jt.random((2,2,3,3))
|
|
b = jt.random((2,2,3))
|
|
ans = jt.linalg.solve(a,b)
|
|
ta = check_solve(a.data,b.data)
|
|
assert np.allclose(ans.data, ta)
|
|
jx = jt.grad(ans, a)
|
|
check_sgrad = jacobian(check_solve)
|
|
gx = check_sgrad(a.data,b.data)
|
|
gx = np.sum(gx,0)
|
|
gx = np.sum(gx,0)
|
|
gx = np.sum(gx,0)
|
|
try:
|
|
assert np.allclose(gx, jx.data,rtol=1)
|
|
except AssertionError:
|
|
print(gx)
|
|
print(jx.data)
|
|
|
|
def test_det(self):
|
|
def check_det(a):
|
|
de = anp.linalg.det(a)
|
|
return de
|
|
|
|
for i in range(50):
|
|
tn = np.random.randn(3, 3).astype('float32') * 5
|
|
while np.allclose(np.linalg.det(tn), 0):
|
|
tn = np.random.randn((3, 3)).astype('float32') * 5
|
|
x = jt.array(tn)
|
|
x = x.reindex([2, 2, x.shape[0], x.shape[1]], ["i2", "i3"])
|
|
s = list(x.shape)
|
|
x_s = s[:-2]
|
|
if len(s) == 2:
|
|
x_s.append(1)
|
|
det = jt.linalg.det(x)
|
|
ta = check_det(x.data)
|
|
assert np.allclose(det.data, ta)
|
|
jx = jt.grad(det, x)
|
|
check_sgrad = jacobian(check_det)
|
|
gx = check_sgrad(x.data)
|
|
gx = np.sum(gx, 2)
|
|
gx = np.sum(gx, 2)
|
|
assert np.allclose(gx, jx.data)
|
|
|
|
|
|
if __name__ == "__main__":
|
|
unittest.main()
|
|
|