def dmrg_shared_exec(mpo_tensors,
init_mps_tensors,
max_mps_rank,
num_iter=1,
sequence='R'):
"""
Perform DMRG iterations with shared executions.
"""
if sequence != "R":
raise NotImplementedError
num = len(mpo_tensors)
size = mpo_tensors[0].shape[1]
mpo_ranks = [mpo_tensors[i].shape[0] for i in range(1, len(mpo_tensors))]
mps_tensors = copy.deepcopy(init_mps_tensors)
mps_ranks = [mps_tensors[i].shape[0] for i in range(1, len(mps_tensors))]
dg = DmrgGraph.create(num, mpo_ranks, mps_ranks, size)
for i, hes in enumerate(dg.hessians):
dg.hessians[i] = simplify(hes)
assert isinstance(hes, ad.EinsumNode)
dg.hessians = generate_sequential_optimal_tree(dg.hessians, dg.mps_inputs)
executor = ad.Executor(dg.hessians)
# sequence is R
for iter in range(num_iter):
mps_tensors = gauge_transform_mps(mps_tensors, right=True)
mps_ranks = [
mps_tensors[i].shape[0] for i in range(1, len(mps_tensors))
]
for i in range(num - 1):
dg.update_graph(num, mpo_ranks, mps_ranks, size)
feed_dict = dict(zip(dg.mpo_inputs, mpo_tensors))
feed_dict.update(dict(zip(dg.mps_inputs, mps_tensors)))
hes_val, = executor.run(feed_dict=feed_dict,
out_nodes=[dg.hessians[i]])
# get the smallest eigenvalue and the corresponding eigenvector of the hesval
eigvec_shape = dg.intermediates[i].shape
eig_val, eigvec = get_smallest_eigenpair(hes_val,
dg.intermediates[i].shape)
# Update the two sites of mps
mps_tensors[i], mps_tensors[i + 1] = dmrg_local_update(
dg.intermediates[i], eigvec, max_mps_rank)
# update the rank
mps_ranks[i] = mps_tensors[i + 1].shape[0]
print(f'At iteration {iter} the smallest eigenvalue is: {eig_val}')
return mps_tensors, eig_val
def cpd_als_shared_exec(dim, size, rank, num_iter, input_val=[]):
A_list, input_tensor, loss, residual = cpd_graph(dim, size, rank)
full_hessian = ad.hessian(loss, A_list)
hessians = [full_hessian[i][i] for i in range(len(full_hessian))]
grads = ad.gradients(loss, A_list)
updates = [
ad.tensordot(ad.tensorinv(hes), grad, [[2, 3], [0, 1]])
for (hes, grad) in zip(hessians, grads)
]
new_A_list = [simplify(A - update) for (A, update) in zip(A_list, updates)]
new_A_list = generate_sequential_optimal_tree(new_A_list, A_list)
executor = ad.Executor(new_A_list)
executor_loss = ad.Executor([simplify(loss)])
if input_val == []:
A_val_list, input_tensor_val = init_rand_cp(dim, size, rank)
else:
A_val_list, input_tensor_val = input_val
for iter in range(num_iter):
t0 = time.time()
# als iterations
for i in range(len(A_list)):
feed_dict = dict(zip(A_list, A_val_list))
feed_dict.update({input_tensor: input_tensor_val})
if i == 0:
A_val_list[0], = executor.run(feed_dict=feed_dict,
out_nodes=[new_A_list[0]])
else:
A_val_list[i], = executor.run(feed_dict=feed_dict,
reset_graph=False,
evicted_inputs=[A_list[i - 1]],
out_nodes=[new_A_list[i]])
feed_dict = dict(zip(A_list, A_val_list))
feed_dict.update({input_tensor: input_tensor_val})
loss_val, = executor_loss.run(feed_dict=feed_dict)
print(f'At iteration {iter} the loss is: {loss_val}')
t1 = time.time()
print(f"[ {iter} ] Sweep took {t1 - t0} seconds")
return A_val_list
def test_dimension_tree_w_identity():
A = ad.Variable(name="A", shape=[2, 2])
B = ad.identity(2)
C = ad.Variable(name="C", shape=[2, 2])
X = ad.Variable(name="X", shape=[2, 2, 2])
einsum_node_A = ad.einsum("abc,bm,cm->am", X, B, C)
einsum_node_B = ad.einsum("abc,am,cm->bm", X, A, C)
einsum_node_C = ad.einsum("abc,am,bm->cm", X, A, B)
dt = generate_sequential_optimal_tree(
[einsum_node_A, einsum_node_B, einsum_node_C], [A, B, C])
assert tree_eq(dt[0], einsum_node_A, [A, C, X])
assert tree_eq(dt[1], einsum_node_B, [A, C, X])
assert tree_eq(dt[2], einsum_node_C, [A, B, X])
def test_dmrg_shared_exec_graph():
from graph_ops.graph_transformer import simplify
from graph_ops.graph_als_optimizer import generate_sequential_optimal_tree
from utils import find_topo_sort
num, rank, size = 4, 3, 2
mpo_ranks = [rank for i in range(1, num)]
mps_ranks = [rank for i in range(1, num)]
dg = DmrgGraph.create(num, mpo_ranks, mps_ranks, size)
for i, hes in enumerate(dg.hessians):
dg.hessians[i] = simplify(hes)
assert isinstance(hes, ad.EinsumNode)
dg.hessians = generate_sequential_optimal_tree(dg.hessians, dg.mps_inputs)
# 8 input variables (4 H term in MPO, 4 A term in MPS), 7 einsum nodes
assert len(find_topo_sort(dg.hessians)) == 15
def tucker_als_graph_shared_exec(dim, size, rank):
"""
Build the graph used for Tucker ALS with shared execution.
Parameters
----------
dim: dimensionality of the input tensor
size: the size of input tensor's each dim
rank: the rank of the decomposition
Returns
-------
tg: an TuckerGraph object
executor: An shared executor
loss: the optimized graph for tucker loss
updates: an list containing updates graphs for each dimension
intermediates: list of einsum nodes. Each node is the objective
each Tucker ALS step optimized for
"""
tg = TuckerGraph(dim, size, rank)
updates = []
for i in range(dim):
core_A = tg.intermediates[i]
hes = ad.hessian(tg.losses[i], [core_A])
hes = hes[0][0]
grad, = ad.gradients(tg.losses[i], [core_A])
new_core_A = core_A - ad.tensordot(
ad.tensorinv(hes), grad,
[[i + dim for i in range(dim)], [i for i in range(dim)]])
updates.append(simplify(new_core_A))
loss = simplify(tg.losses[0])
for i in range(1, len(tg.losses)):
assert loss.name == simplify(tg.losses[i]).name
updates = generate_sequential_optimal_tree(updates, tg.A_list)
executor_updates = ad.Executor(updates)
executor_loss = ad.Executor([loss])
return tg, executor_updates, executor_loss, loss, updates, tg.intermediates
def test_simple_dmrg_tree():
A1 = ad.Variable(name="A1", shape=[3, 2])
A2 = ad.Variable(name="A2", shape=[3, 3, 2])
A3 = ad.Variable(name="A3", shape=[3, 2])
X1 = ad.Variable(name="X1", shape=[3, 2, 2])
X2 = ad.Variable(name="X2", shape=[3, 3, 2, 2])
X3 = ad.Variable(name="X3", shape=[3, 2, 2])
"""
The network and indices positions are as follows:
A1 - f - A2 - g - A3
| | |
c d e
| | |
X1 - a - X2 - b - X3
| | |
h i j
| | |
A1 - k - A2 - l - A3
"""
einsum_node_A1 = ad.einsum("ach,abdi,bej,fgd,kli,ge,lj->fckh", X1, X2, X3,
A2, A2, A3, A3)
einsum_node_A2 = ad.einsum("ach,abdi,bej,fc,kh,ge,lj->fgdkli", X1, X2, X3,
A1, A1, A3, A3)
einsum_node_A3 = ad.einsum("ach,abdi,bej,fc,kh,fgd,kli->gelj", X1, X2, X3,
A1, A1, A2, A2)
dt = generate_sequential_optimal_tree(
[einsum_node_A1, einsum_node_A2, einsum_node_A3], [A1, A2, A3])
assert tree_eq(dt[0], einsum_node_A1, [X1, X2, X3, A1, A1, A2, A2, A3, A3])
assert tree_eq(dt[1], einsum_node_A2, [X1, X2, X3, A1, A1, A2, A2, A3, A3])
# In the correct contraction path, only X3 should be contracted with A3,
# all other X nodes should be contracted later.
einsum_inputs = list(
filter(lambda node: isinstance(node, ad.EinsumNode),
find_topo_sort(dt)))
assert sorted(einsum_inputs[0].inputs,
key=lambda node: node.name) == sorted(
[A3, A3, X3], key=lambda node: node.name)
def test_dimension_tree_4d():
A = ad.Variable(name="A", shape=[2, 2])
B = ad.Variable(name="B", shape=[2, 2])
C = ad.Variable(name="C", shape=[2, 2])
D = ad.Variable(name="D", shape=[2, 2])
X = ad.Variable(name="X", shape=[2, 2, 2, 2])
einsum_node_A = ad.einsum("abcd,bm,cm,dm->am", X, B, C, D)
einsum_node_B = ad.einsum("abcd,am,cm,dm->bm", X, A, C, D)
einsum_node_C = ad.einsum("abcd,am,bm,dm->cm", X, A, B, D)
einsum_node_D = ad.einsum("abcd,am,bm,cm->dm", X, A, B, C)
dt = generate_sequential_optimal_tree(
[einsum_node_A, einsum_node_B, einsum_node_C, einsum_node_D],
[A, B, C, D])
# 5 inputs, 4 outputs, 5 intermedaites
assert len(find_topo_sort(dt)) == 14
assert tree_eq(dt[0], einsum_node_A, [A, B, C, D, X])
assert tree_eq(dt[1], einsum_node_B, [A, B, C, D, X])
assert tree_eq(dt[2], einsum_node_C, [A, B, C, D, X])
assert tree_eq(dt[3], einsum_node_D, [A, B, C, D, X])
def dmrg_shared_exec_iterative_solve(mpo_tensors,
init_mps_tensors,
max_mps_rank,
num_iter=1,
sequence='R'):
"""
Perform DMRG iterations with shared execution and iterative solve.
"""
if sequence != "R":
raise NotImplementedError
num = len(mpo_tensors)
size = mpo_tensors[0].shape[1]
mpo_ranks = [mpo_tensors[i].shape[0] for i in range(1, len(mpo_tensors))]
mps_tensors = copy.deepcopy(init_mps_tensors)
mps_ranks = [mps_tensors[i].shape[0] for i in range(1, len(mps_tensors))]
dg = DmrgImplicitUpdateGraph.create(num, mpo_ranks, mps_ranks, size)
for i, hvp in enumerate(dg.hvps):
dg.hvps[i] = simplify(hvp)
assert isinstance(hvp, ad.EinsumNode)
dg.hvps = generate_sequential_optimal_tree(dg.hvps, dg.mps_inputs)
executor_hvps = ad.Executor(dg.hvps)
executor_intermediates = ad.Executor(dg.intermediates)
# sequence is R
for iter in range(num_iter):
mps_tensors = gauge_transform_mps(mps_tensors, right=True)
mps_ranks = [
mps_tensors[i].shape[0] for i in range(1, len(mps_tensors))
]
for i in range(num - 1):
dg.update_graph(num, mpo_ranks, mps_ranks, size)
feed_dict = dict(zip(dg.mpo_inputs, mpo_tensors))
feed_dict.update(dict(zip(dg.mps_inputs, mps_tensors)))
intermediate, = executor_intermediates.run(
feed_dict=feed_dict, out_nodes=[dg.intermediates[i]])
# Calculate the eigenvector using the implicit solver.
# Note: This only supports NumPy datatype.
# TODO: Add a general Lanczos solver that adapts to all the backends.
operator = DMRGLinearOperator(dg, executor_hvps, i, feed_dict)
# Reference: https://docs.scipy.org/doc/scipy/reference/generated/scipy.sparse.linalg.eigsh.html
eig_vals, eigvecs = spla.eigsh(operator,
k=1,
ncv=4,
tol=1e-3,
which='SA',
v0=intermediate.ravel())
eig_val, eigvec = eig_vals[0], eigvecs[:, 0]
eigvec = T.reshape(eigvec, dg.intermediates[i].shape)
# Update the two sites of mps
mps_tensors[i], mps_tensors[i + 1] = dmrg_local_update(
dg.intermediates[i], eigvec, max_mps_rank)
# update the rank
mps_ranks[i] = mps_tensors[i + 1].shape[0]
print(f'At site {i}, the smallest eigenvalue is: {eig_val}')
print(f'At iteration {iter} the smallest eigenvalue is: {eig_val}')
return mps_tensors, eig_val