def testTaylorExpansion(self, shape):
def f_2_exact(x0, x, params, do_alter, do_shift_x=True):
w1, w2, b = params
f_lin = EmpiricalTest.f_lin_exact(x0, x, params, do_alter,
do_shift_x)
if do_shift_x:
x0 = x0 * 2 + 1.
x = x * 2 + 1.
if do_alter:
b *= 2.
w1 += 5.
w2 /= 0.9
dx = x - x0
return f_lin + 0.5 * np.dot(np.dot(dx.T, w1), dx)
key = random.PRNGKey(0)
key, s1, s2, s3, = random.split(key, 4)
w1 = random.normal(s1, shape)
w1 = 0.5 * (w1 + w1.T)
w2 = random.normal(s2, shape)
b = random.normal(s3, (shape[-1], ))
params = (w1, w2, b)
key, split = random.split(key)
x0 = random.normal(split, (shape[-1], ))
f_lin = empirical.taylor_expand(EmpiricalTest.f, x0, 1)
f_2 = empirical.taylor_expand(EmpiricalTest.f, x0, 2)
for _ in range(TAYLOR_RANDOM_SAMPLES):
for do_alter in [True, False]:
for do_shift_x in [True, False]:
key, split = random.split(key)
x = random.normal(split, (shape[-1], ))
self.assertAllClose(
EmpiricalTest.f_lin_exact(x0,
x,
params,
do_alter,
do_shift_x=do_shift_x),
f_lin(x, params, do_alter, do_shift_x=do_shift_x),
True)
self.assertAllClose(
f_2_exact(x0,
x,
params,
do_alter,
do_shift_x=do_shift_x),
f_2(x, params, do_alter, do_shift_x=do_shift_x), True)
def testTaylorExpansion(self, shape):
# We use a three layer deep linear network for testing.
def f(x, params):
w1, w2, b = params
return 0.5 * np.dot(np.dot(x.T, w1), x) + np.dot(w2, x) + b
def f_lin_exact(x0, x, params):
w1, w2, b = params
f0 = f(x0, params)
dx = x - x0
return f0 + np.dot(np.dot(x0.T, w1) + w2, dx)
def f_2_exact(x0, x, params):
w1, w2, b = params
dx = x - x0
return f_lin_exact(x0, x,
params) + 0.5 * np.dot(np.dot(dx.T, w1), dx)
key = random.PRNGKey(0)
key, s1, s2, s3, = random.split(key, 4)
w1 = random.normal(s1, shape)
w1 = 0.5 * (w1 + w1.T)
w2 = random.normal(s2, shape)
b = random.normal(s3, (shape[-1], ))
params = (w1, w2, b)
key, split = random.split(key)
x0 = random.normal(split, (shape[-1], ))
f_lin = empirical.taylor_expand(f, x0, 1)
f_2 = empirical.taylor_expand(f, x0, 2)
for _ in range(TAYLOR_RANDOM_SAMPLES):
key, split = random.split(key)
x = random.normal(split, (shape[-1], ))
self.assertAllClose(f_lin_exact(x0, x, params), f_lin(x, params),
True)
self.assertAllClose(f_2_exact(x0, x, params), f_2(x, params), True)
def testTaylorExpansion(self, shape):
def f_2_exact(x0, x, params, do_alter, do_shift_x=True):
w1, w2, b = params
f_lin = EmpiricalTest.f_lin_exact(x0, x, params, do_alter,
do_shift_x)
if do_shift_x:
x0 = x0 * 2 + 1.
x = x * 2 + 1.
if do_alter:
b *= 2.
w1 += 5.
w2 /= 0.9
dx = x - x0
return f_lin + 0.5 * np.dot(np.dot(dx.T, w1), dx)
key = stateless_uniform(shape=[2],
seed=[0, 0],
minval=None,
maxval=None,
dtype=tf.int32)
splits = tf_random_split(seed=tf.convert_to_tensor(key,
dtype=tf.int32),
num=4)
key = splits[0]
s1 = splits[1]
s2 = splits[2]
s3 = splits[3]
w1 = np.asarray(normal(shape, seed=s1))
w1 = 0.5 * (w1 + w1.T)
w2 = np.asarray(normal(shape, seed=s2))
b = np.asarray(normal((shape[-1], ), seed=s3))
params = (w1, w2, b)
splits = tf_random_split(seed=tf.convert_to_tensor(key,
dtype=tf.int32),
num=2)
key = splits[0]
split = splits[1]
x0 = np.asarray(normal((shape[-1], ), seed=split))
f_lin = empirical.taylor_expand(EmpiricalTest.f, x0, 1)
f_2 = empirical.taylor_expand(EmpiricalTest.f, x0, 2)
for _ in range(TAYLOR_RANDOM_SAMPLES):
for do_alter in [True, False]:
for do_shift_x in [True, False]:
splits = tf_random_split(seed=tf.convert_to_tensor(
key, dtype=tf.int32),
num=2)
key = splits[0]
split = splits[1]
x = np.asarray(normal((shape[-1], ), seed=split))
self.assertAllClose(
EmpiricalTest.f_lin_exact(x0,
x,
params,
do_alter,
do_shift_x=do_shift_x),
f_lin(x, params, do_alter, do_shift_x=do_shift_x))
self.assertAllClose(
f_2_exact(x0,
x,
params,
do_alter,
do_shift_x=do_shift_x),
f_2(x, params, do_alter, do_shift_x=do_shift_x))