def contract_network(net: batchtensornetwork.BatchTensorNetwork,
var_nodes: List[tensornetwork.Node],
data_nodes: Tuple[tensornetwork.Node]) -> tf.Tensor:
"""Contracts TensorNetwork created in `create_network`."""
batch_edges = tuple((x[0] for x in data_nodes))
space_edge = data_nodes[1][1]
# Contract data with the MPS tensors
var_nodes[0] = net.contract_between(data_nodes[0], var_nodes[0])
var_nodes[1] = net.batched_contract_between(data_nodes[1],
var_nodes[1], space_edge,
var_nodes[1][0])
# Contract the artificial "space" index. This step is equivalent to
# contracting the MPS over the bond dimensions.
var_nodes[1] = batchtensornetwork.pairwise_reduction(
net, var_nodes[1], space_edge)
# Contract the final bond edge with the boundary.
var_nodes = net.batched_contract_between(var_nodes[0], var_nodes[1],
batch_edges[0],
batch_edges[1])
return var_nodes.tensor
def test_pairwise_reduction_multiple_edges():
net = batchtensornetwork.BatchTensorNetwork(backend="tensorflow")
a = net.add_node(np.ones([3, 5, 6, 2, 2]))
b = batchtensornetwork.pairwise_reduction(net, a, a[1])
np.testing.assert_allclose(b.tensor, 16 * np.ones([3, 6, 2, 2]))
def test_pairwise_reduction():
net = batchtensornetwork.BatchTensorNetwork(backend="tensorflow")
a = net.add_node(np.ones([10, 2, 2]))
b = batchtensornetwork.pairwise_reduction(net, a, a[0])
np.testing.assert_allclose(b.tensor, 512 * np.ones([2, 2]))