def test_tensor_jacobian_product(): # This function will have an asymmetric jacobian matrix. fun = lambda a: np.roll(np.sin(a), 1) a = npr.randn(5) V = npr.randn(5) J = jacobian(fun)(a) check_equivalent(np.dot(V.T, J), tensor_jacobian_product(fun)(a, V))
def test_hvp(): fun = lambda a: np.sum(np.sin(a)) a = npr.randn(5) v = npr.randn(5) H = hessian(fun)(a) hvp = make_hvp(fun)(a)[0] check_equivalent(np.dot(H, v), hvp(v))
def test_hessian_tensor_product(): fun = lambda a: np.sum(np.sin(a)) a = npr.randn(5, 4, 3) V = npr.randn(5, 4, 3) H = hessian(fun)(a) check_equivalent(np.tensordot(H, V, axes=np.ndim(V)), hessian_tensor_product(fun)(a, V))
def test_tensor_jacobian_product(): fun = lambda a: np.roll(np.sin(a), 1) a = npr.randn(5, 4, 3) V = npr.randn(5, 4) J = jacobian(fun)(a) check_equivalent(np.tensordot(V, J, axes=np.ndim(V)), tensor_jacobian_product(fun)(a, V))
def test_hessian(): # Check Hessian of a quadratic function. D = 5 H = npr.randn(D, D) def fun(x): return np.dot(np.dot(x, H),x) hess = hessian(fun) x = npr.randn(D) check_equivalent(hess(x), H + H.T)
def test_elementwise_grad(): def simple_fun(a): return a + np.sin(a) + np.cosh(a) A = npr.randn(10) wrapped = elementwise_grad(simple_fun)(A) explicit = np.array([grad(simple_fun)(A[i]) for i in range(len(A))]) check_equivalent(wrapped, explicit)
def test_make_ggnvp(): A = npr.randn(5, 4) x = npr.randn(4) v = npr.randn(4) fun = lambda x: np.dot(A, x) check_equivalent(make_ggnvp(fun)(x)(v), _make_explicit_ggnvp(fun)(x)(v)) fun2 = lambda x: np.tanh(np.dot(A, x)) check_equivalent(make_ggnvp(fun2)(x)(v), _make_explicit_ggnvp(fun2)(x)(v))
def test_make_jvp(): A = npr.randn(3, 5) x = npr.randn(5) v = npr.randn(5) fun = lambda x: np.tanh(np.dot(A, x)) jvp_explicit = lambda x: lambda v: np.dot(jacobian(fun)(x), v) jvp = make_jvp(fun) check_equivalent(jvp_explicit(x)(v), jvp(x)(v)[1])
def test_grad_and_aux(): A = npr.randn(5, 4) x = npr.randn(4) f = lambda x: (np.sum(np.dot(A, x)), x**2) g = lambda x: np.sum(np.dot(A, x)) assert len(grad_and_aux(f)(x)) == 2 check_equivalent(grad_and_aux(f)(x)[0], grad(g)(x)) check_equivalent(grad_and_aux(f)(x)[1], x**2)
def test_value_and_grad(): fun = lambda x: np.sum(np.sin(x)**2) dfun = grad(fun) dfun_both = value_and_grad(fun) x = npr.randn(5) assert not isbox(dfun_both(x)[0]) check_equivalent(fun(x), dfun_both(x)[0]) check_equivalent(dfun(x), dfun_both(x)[1]) def fun2(x): return dfun_both(x)[0] check_grads(fun2)(x)
def test_hessian(): # Check Hessian of a quadratic function. D = 5 H = npr.randn(D, D) def fun(x): return np.dot(np.dot(x, H), x) hess = hessian(fun) x = npr.randn(D) check_equivalent(hess(x), H + H.T)
def test_elementwise_grad_multiple_args(): def simple_fun(a, b): return a + np.sin(a) + np.cosh(b) A = 0.9 B = npr.randn(10) argnum = 1 wrapped = elementwise_grad(simple_fun, argnum)(A, B) explicit = np.array([grad(simple_fun, argnum)(A, B[i]) for i in range(len(B))]) check_equivalent(wrapped, explicit)
def test_make_ggnvp_nondefault_g(): A = npr.randn(5, 4) x = npr.randn(4) v = npr.randn(4) g = lambda y: np.sum(2.*y**2 + y**4) fun = lambda x: np.dot(A, x) check_equivalent(make_ggnvp(fun, g)(x)(v), _make_explicit_ggnvp(fun, g)(x)(v)) fun2 = lambda x: np.tanh(np.dot(A, x)) check_equivalent(make_ggnvp(fun2, g)(x)(v), _make_explicit_ggnvp(fun2, g)(x)(v))
def test_elementwise_grad_multiple_args(): def simple_fun(a, b): return a + np.sin(a) + np.cosh(b) A = 0.9 B = npr.randn(10) argnum = 1 wrapped = elementwise_grad(simple_fun, argnum)(A, B) explicit = np.array( [grad(simple_fun, argnum)(A, B[i]) for i in range(len(B))]) check_equivalent(wrapped, explicit)
def test_make_ggnvp_nondefault_g(): A = npr.randn(5, 4) x = npr.randn(4) v = npr.randn(4) g = lambda y: np.sum(2. * y**2 + y**4) fun = lambda x: np.dot(A, x) check_equivalent( make_ggnvp(fun, g)(x)(v), _make_explicit_ggnvp(fun, g)(x)(v)) fun2 = lambda x: np.tanh(np.dot(A, x)) check_equivalent( make_ggnvp(fun2, g)(x)(v), _make_explicit_ggnvp(fun2, g)(x)(v))
def test_multigrad(): def complicated_fun(a,b,c,d,e,f=1.1, g=9.0): return a + np.sin(b) + np.cosh(c) + np.cos(d) + np.tan(e) + f + g def complicated_fun_3_1(d_b): d, b = d_b return complicated_fun(A, b, C, d, E, f=F, g=G) A = 0.5 B = -0.3 C = 0.2 D = -1.1 E = 0.7 F = 0.6 G = -0.1 wrapped = grad(complicated_fun, argnum=[3, 1])(A, B, C, D, E, f=F, g=G) explicit = grad(complicated_fun_3_1)((D, B)) check_equivalent(wrapped, explicit)
def test_multigrad(): def complicated_fun(a, b, c, d, e, f=1.1, g=9.0): return a + np.sin(b) + np.cosh(c) + np.cos(d) + np.tan(e) + f + g def complicated_fun_3_1(d_b): d, b = d_b return complicated_fun(A, b, C, d, E, f=F, g=G) A = 0.5 B = -0.3 C = 0.2 D = -1.1 E = 0.7 F = 0.6 G = -0.1 wrapped = grad(complicated_fun, argnum=[3, 1])(A, B, C, D, E, f=F, g=G) explicit = grad(complicated_fun_3_1)((D, B)) check_equivalent(wrapped, explicit)
def test_value_and_multigrad(): def complicated_fun(a, b, c, d, e, f=1.1, g=9.0): return a + np.sin(b) + np.cosh(c) + np.cos(d) + np.tan(e) + f + g A = 0.5 B = -0.3 C = 0.2 D = -1.1 E = 0.7 F = 0.6 G = -0.1 dfun = grad(complicated_fun, argnum=[3, 1]) dfun_both = value_and_grad(complicated_fun, argnum=[3, 1]) check_equivalent(complicated_fun(A, B, C, D, E, f=F, g=G), dfun_both(A, B, C, D, E, f=F, g=G)[0]) check_equivalent(dfun(A, B, C, D, E, f=F, g=G), dfun_both(A, B, C, D, E, f=F, g=G)[1])
def test_value_and_multigrad(): def complicated_fun(a,b,c,d,e,f=1.1, g=9.0): return a + np.sin(b) + np.cosh(c) + np.cos(d) + np.tan(e) + f + g A = 0.5 B = -0.3 C = 0.2 D = -1.1 E = 0.7 F = 0.6 G = -0.1 dfun = grad(complicated_fun, argnum=[3, 1]) dfun_both = value_and_grad(complicated_fun, argnum=[3, 1]) check_equivalent(complicated_fun(A, B, C, D, E, f=F, g=G), dfun_both(A, B, C, D, E, f=F, g=G)[0]) check_equivalent(dfun(A, B, C, D, E, f=F, g=G), dfun_both(A, B, C, D, E, f=F, g=G)[1])
def test_matrix_jacobian_product(): fun = lambda a: np.roll(np.sin(a), 1) a = npr.randn(5, 4) V = npr.randn(5, 4) J = jacobian(fun)(a) check_equivalent(np.tensordot(V, J), tensor_jacobian_product(fun)(a, V))
def test_hessian_matrix_product(): fun = lambda a: np.sum(np.sin(a)) a = npr.randn(5, 4) V = npr.randn(5, 4) H = hessian(fun)(a) check_equivalent(np.tensordot(H, V), hessian_tensor_product(fun)(a, V))
def test_hessian_tensor_product(): fun = lambda a: np.sum(np.sin(a)) a = npr.randn(5) v = npr.randn(5) H = hessian(fun)(a) check_equivalent(np.dot(H, v), hessian_tensor_product(fun)(a, v))
def test_multigrad_onearg(): fun = lambda x, y: np.sum(x + np.sin(y)) packed_fun = lambda xy: np.sum(xy[0] + np.sin(xy[1])) A, B = npr.randn(3), npr.randn(3) check_equivalent(grad(fun, argnum=[0])(A,B), (grad(packed_fun)((A,B))[0],))
def test_multigrad_onearg(): fun = lambda x, y: np.sum(x + np.sin(y)) packed_fun = lambda xy: np.sum(xy[0] + np.sin(xy[1])) A, B = npr.randn(3), npr.randn(3) check_equivalent( grad(fun, argnum=[0])(A, B), (grad(packed_fun)((A, B))[0], ))