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