def test_s2s_jtjvp(backendopt):
for datatype in backendopt:
T.set_backend(datatype)
x = ad.Variable(name="x", shape=[2])
A = ad.Variable(name="A", shape=[3, 2])
v = ad.Variable(name="v", shape=[2])
y = ad.einsum("ab,b->a", A, x)
jtjvp_x, = ad.jtjvps(y, [x], [v])
x_val = T.tensor([1., 2.])
A_val = T.tensor([[1., 2.], [3., 4.], [5, 6]])
v_val = T.tensor([3, 4])
expected_jtjvp_x_val = T.einsum("ba,bc,c->a", A_val, A_val, v_val)
StS = SourceToSource()
forward_str = StS.forward([jtjvp_x],
function_name='jtjvp',
backend=datatype)
m = import_code(forward_str)
jtjvp_x_val_s2s, = m.jtjvp([A_val, v_val])
assert isinstance(jtjvp_x, ad.Node)
assert T.array_equal(jtjvp_x_val_s2s, expected_jtjvp_x_val)
def test_jtjvps(backendopt):
for datatype in backendopt:
T.set_backend(datatype)
x = ad.Variable(name="x", shape=[2])
A = ad.Variable(name="A", shape=[3, 2])
v = ad.Variable(name="v", shape=[2])
y = ad.einsum('ab, b->a', A, x)
jtjvp_x, = ad.jtjvps(y, [x], [v])
executor = ad.Executor([y, jtjvp_x])
x_val = T.tensor([1., 2.])
A_val = T.tensor([[1., 2.], [3., 4.], [5, 6]])
v_val = T.tensor([3., 4.])
y_val, jtjvp_x_val = executor.run(feed_dict={
x: x_val,
A: A_val,
v: v_val
})
expected_yval = T.einsum('ab, b->a', A_val, x_val)
expected_jtjvp_x_val = T.einsum('ba, ac->bc', T.transpose(A_val),
A_val)
expected_jtjvp_x_val = T.einsum('ab, b->a', expected_jtjvp_x_val,
v_val)
assert isinstance(jtjvp_x, ad.Node)
assert T.array_equal(y_val, expected_yval)
assert T.array_equal(jtjvp_x_val, expected_jtjvp_x_val)
def test_cpd_jtjvp_optimize(backendopt):
dim = 3
for datatype in backendopt:
T.set_backend(datatype)
A_list, input_tensor, loss, residual = cpd_graph(dim, size, rank)
A, B, C = A_list
v_A = ad.Variable(name="v_A", shape=[size, rank])
v_B = ad.Variable(name="v_B", shape=[size, rank])
v_C = ad.Variable(name="v_C", shape=[size, rank])
A_list, input_tensor_val = init_rand_cp(dim, size, rank)
A_val, B_val, C_val = A_list
v_A_list, _ = init_rand_cp(dim, size, rank)
v_A_val, v_B_val, v_C_val = v_A_list
JtJvps = ad.jtjvps(output_node=residual,
node_list=[A, B, C],
vector_list=[v_A, v_B, v_C])
JtJvps = [optimize(JtJvp) for JtJvp in JtJvps]
dedup(*JtJvps)
for node in JtJvps:
assert isinstance(node, ad.AddNode)
executor_JtJvps = ad.Executor(JtJvps)
jtjvp_val = executor_JtJvps.run(
feed_dict={
A: A_val,
B: B_val,
C: C_val,
input_tensor: input_tensor_val,
v_A: v_A_val,
v_B: v_B_val,
v_C: v_C_val
})
expected_hvp_val = expect_jtjvp_val(A_val, B_val, C_val, v_A_val,
v_B_val, v_C_val)
assert T.norm(jtjvp_val[0] - expected_hvp_val[0]) < 1e-8
assert T.norm(jtjvp_val[1] - expected_hvp_val[1]) < 1e-8
assert T.norm(jtjvp_val[2] - expected_hvp_val[2]) < 1e-8
def cpd_nls(size, rank, regularization=1e-7, mode='ad'):
"""
mode: ad / optimized / jax
"""
assert mode in {'ad', 'jax', 'optimized'}
dim = 3
for datatype in BACKEND_TYPES:
T.set_backend(datatype)
T.seed(1)
A_list, input_tensor, loss, residual = cpd_graph(dim, size, rank)
A, B, C = A_list
v_A = ad.Variable(name="v_A", shape=[size, rank])
v_B = ad.Variable(name="v_B", shape=[size, rank])
v_C = ad.Variable(name="v_C", shape=[size, rank])
grads = ad.gradients(loss, [A, B, C])
JtJvps = ad.jtjvps(output_node=residual,
node_list=[A, B, C],
vector_list=[v_A, v_B, v_C])
A_list, input_tensor_val = init_rand_cp(dim, size, rank)
A_val, B_val, C_val = A_list
if mode == 'jax':
from source import SourceToSource
StS = SourceToSource()
StS.forward(JtJvps,
file=open("examples/jax_jtjvp.py", "w"),
function_name='jtjvp',
backend='jax')
executor_grads = ad.Executor([loss] + grads)
JtJvps = [optimize(JtJvp) for JtJvp in JtJvps]
dedup(*JtJvps)
executor_JtJvps = ad.Executor(JtJvps)
optimizer = cp_nls_optimizer(input_tensor_val, [A_val, B_val, C_val])
regu_increase = False
normT = T.norm(input_tensor_val)
time_all, fitness = 0., 0.
for i in range(10):
t0 = time.time()
def hess_fn(v):
if mode == 'optimized':
from examples.cpd_jtjvp_optimized import jtjvp
return jtjvp([v[0], B_val, C_val, v[1], A_val, v[2]])
elif mode == 'ad':
return executor_JtJvps.run(
feed_dict={
A: A_val,
B: B_val,
C: C_val,
input_tensor: input_tensor_val,
v_A: v[0],
v_B: v[1],
v_C: v[2]
})
elif mode == 'jax':
from examples.jax_jtjvp import jtjvp
return jtjvp([B_val, C_val, v[0], A_val, v[1], v[2]])
loss_val, grad_A_val, grad_B_val, grad_C_val = executor_grads.run(
feed_dict={
A: A_val,
B: B_val,
C: C_val,
input_tensor: input_tensor_val
})
res = math.sqrt(loss_val)
fitness = 1 - res / normT
print(f"[ {i} ] Residual is {res} fitness is: {fitness}")
print(f"Regularization is: {regularization}")
[A_val, B_val, C_val], total_cg_time = optimizer.step(
hess_fn=hess_fn,
grads=[grad_A_val / 2, grad_B_val / 2, grad_C_val / 2],
regularization=regularization)
t1 = time.time()
print(f"[ {i} ] Sweep took {t1 - t0} seconds")
time_all += t1 - t0
if regularization < 1e-07:
regu_increase = True
elif regularization > 1:
regu_increase = False
if regu_increase:
regularization = regularization * 2
else:
regularization = regularization / 2
return total_cg_time, fitness