def tangent_func():
func.__globals__['np'] = np
df = tangent.autodiff(func,
mode='forward',
preserve_result=preserve_result,
wrt=wrt,
optimized=True,
verbose=1)
args_ = args + (1.0, ) # seed gradient
return df(*deepcopy(args_))
def _test_tf_hvp(func, optimized):
a = tf.random_normal(shape=(300,))
v = tf.reshape(a, shape=(-1,))
modes = ['forward', 'reverse']
for mode1 in modes:
for mode2 in modes:
if mode1 == mode2 == 'forward':
continue
df = tangent.autodiff(
func,
mode=mode1,
motion='joint',
optimized=optimized,
check_dims=False)
ddf = tangent.autodiff(
df, mode=mode2, motion='joint', optimized=optimized, check_dims=False)
dx = ddf(a, tf.constant(1.0), v)
# We just ensure it computes something in this case.
assert dx.shape == a.shape
def _test_hvp(func, optimized):
np.random.seed(0)
a = np.random.normal(scale=1, size=(300,)).astype('float32')
v = a.ravel()
modes = ['forward', 'reverse']
for mode1 in modes:
for mode2 in modes:
if mode1 == mode2 == 'forward':
continue
df = tangent.autodiff(
func,
mode=mode1,
motion='joint',
optimized=optimized,
check_dims=False)
ddf = tangent.autodiff(
df, mode=mode2, motion='joint', optimized=optimized, check_dims=False)
dx = ddf(a, 1, v)
hvp_ag = hessian_vector_product(func)
dx_ag = hvp_ag(a, v)
assert np.allclose(dx, dx_ag)
def test_nested_dict(motion, optimized):
p = dict(i=dict(j=3.0, k=4.0))
func = nested_dict
df = tangent.autodiff(
func,
motion=motion,
optimized=optimized,
verbose=True,
input_derivative=INPUT_DERIVATIVE.DefaultOne)
dx = df(p)
df_ag = ag_grad(func)
dx_ag = df_ag(p)
for k in p['i']:
assert np.allclose(dx['i'][k], dx_ag['i'][k])
def tangent_func():
y = func(*deepcopy(args))
if np.array(y).size > 1:
init_grad = np.ones_like(y)
else:
init_grad = 1
func.__globals__['np'] = np
df = tangent.autodiff(func,
mode='reverse',
motion=motion,
optimized=optimized,
preserve_result=preserve_result,
verbose=1)
if motion == 'joint':
return df(*deepcopy(args) + (init_grad, ))
return df(*deepcopy(args), init_grad=init_grad)
def test_unpacking_args_saxpy(motion, optimized, a, b, c):
func = unpacking_args_saxpy
func = tangent.tangent(func)
func.__globals__['np'] = np
df = tangent.autodiff(
func,
motion=motion,
optimized=optimized,
verbose=True,
input_derivative=INPUT_DERIVATIVE.DefaultOne)
dx = df((a, b, c))
df_num = utils.numeric_grad(func)
dx_num = df_num((a, b, c))
assert np.allclose(dx, dx_num)
def test_inlining_contextmanager(motion, optimized, a):
func = inlining_contextmanager
func = tangent.tangent(func)
func.__globals__['np'] = np
df = tangent.autodiff(
func,
motion=motion,
optimized=optimized,
verbose=True,
input_derivative=INPUT_DERIVATIVE.DefaultOne)
dx = df(a)
func.__globals__['np'] = ag_np
df_ag = ag_grad(func)
df_ag(a)
assert np.allclose(dx, 2.9 * a**2)
def test_dict_saxpy(motion, optimized, a, b, c):
func = dict_saxpy
func = tangent.tangent(func)
func.__globals__['np'] = np
df = tangent.autodiff(
func,
motion=motion,
optimized=optimized,
verbose=True,
input_derivative=INPUT_DERIVATIVE.DefaultOne)
dx = df(dict(a=a, b=b, c=c))
df_num = utils.numeric_grad(func)
dx_num = df_num(dict(a=float(a), b=float(b), c=float(c)))
flat_dx, _ = flatten(dx)
flat_dx_num, _ = flatten(dx_num)
assert np.allclose(flat_dx, flat_dx_num)
def make_M(F, B, NV):
""" Returns a function that calculates the system Jacobian, given the flux
function, and the non-conservative matrix (if necessary)
"""
dFdQ = autodiff(F)
def M(Q, d, pars=None):
""" Returns the system Jacobian in direction d, given state Q
"""
ret = zeros([NV, NV])
for i in range(NV):
x = zeros(NV)
x[i] = 1
ret[i] = dFdQ(Q, d, pars, x)
if B is not None:
ret += B(Q, d, pars)
return ret
return M
def test_rnn(motion, optimized):
func = rnn
w = np.random.randn(2, 3)
inputs = np.random.randn(3, 2)
func.__globals__['np'] = np
df = tangent.autodiff(
func,
wrt=(0, 1),
motion=motion,
optimized=optimized,
verbose=True,
input_derivative=INPUT_DERIVATIVE.DefaultOne)
dinputs, dw = df(inputs, w)
num_dinputs = utils.numeric_grad(func)(inputs, w)
num_dw = utils.numeric_grad(lambda w, x: func(x, w))(w, inputs)
assert np.allclose(num_dw, dw)
assert np.allclose(num_dinputs, dinputs)
def test_bilinear(optimized):
func = bilinear
D = 3
np.random.seed(0)
x = np.random.randn(1, D)
h = np.random.randn(1, D)
U = np.random.randn(D, D)
w = np.random.randn(D, D)
b = np.random.randn(1, D)
func.__globals__['np'] = np
df = tangent.autodiff(
func,
wrt=(0,),
motion='joint',
optimized=optimized,
verbose=True,
input_derivative=INPUT_DERIVATIVE.DefaultOne)
dx = df(x, h, U, w, b)
num_dx = utils.numeric_grad(func)(x, h, U, w, b)
assert np.allclose(num_dx, dx)
def test_logistic_regression(motion, optimized):
func = logistic_regression
w = np.random.randn(3, 5)
b = np.random.randn(5)
input_ = np.random.rand(3)
label = np.zeros(5)
label[1] = 1
func.__globals__['np'] = np
df = tangent.autodiff(
func,
wrt=(2, 3),
motion=motion,
optimized=optimized,
verbose=True,
input_derivative=INPUT_DERIVATIVE.DefaultOne)
dw, db = df(input_, label, w, b)
func.__globals__['np'] = ag_np
ag_dw = ag_grad(func, argnum=2)(input_, label, w, b)
ag_db = ag_grad(func, argnum=3)(input_, label, w, b)
assert np.allclose(ag_dw, dw)
assert np.allclose(ag_db, db)
def assert_forward_not_implemented(func, wrt):
try:
tangent.autodiff(func, mode='forward', preserve_result=False, wrt=wrt)
assert False, 'Remove this when implementing.'
except NotImplementedError:
pass
