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_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():
# 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_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_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_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_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_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_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_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_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_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_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_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_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_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], ))
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))
