def call(self, inputs: tf.Tensor, **kwargs) -> tf.Tensor:
tn_nodes = [tn.Node(n, backend='tensorflow') for n in self.nodes]
for i in range(len(tn_nodes) - 1):
tn_nodes[i][2] ^ tn_nodes[i + 1][1]
input_edges = [n[0] for n in tn_nodes]
output_edges = [n[3] for n in tn_nodes]
edges = [tn_nodes[0][1], tn_nodes[-1][2]] + input_edges + output_edges
contracted = tn.contractors.greedy(tn_nodes, edges)
tn.flatten_edges(input_edges)
tn.flatten_edges(output_edges)
tf_df = 'NCHW' if self.data_format == 'channels_first' else 'NHWC'
result = tf.nn.conv2d(inputs,
contracted.tensor,
self.strides,
self.padding.upper(),
data_format=tf_df,
dilations=self.dilation_rate)
if self.use_bias:
bias = tf.reshape(self.bias_var, (
1,
self.filters,
))
result += bias
if self.activation is not None:
result = self.activation(result)
return result
def test_flatten_edges_different_nodes_value_error(backend):
a = tn.Node(np.eye(2), backend=backend)
b = tn.Node(np.eye(2), backend=backend)
c = tn.Node(np.eye(2), backend=backend)
e1 = tn.connect(a[0], b[0])
e2 = tn.connect(a[1], c[0])
tn.connect(b[1], c[1])
with pytest.raises(ValueError):
tn.flatten_edges([e1, e2])
def test_flatten_trace_consistent_tensor(backend):
a_val = np.ones((5, 3, 4, 4, 5))
a = tn.Node(a_val, backend=backend)
e1 = tn.connect(a[0], a[4])
e2 = tn.connect(a[3], a[2])
tn.flatten_edges([e2, e1])
tn.check_correct({a})
# Check expected values.
a_final = np.reshape(np.transpose(a_val, (1, 2, 0, 3, 4)), (3, 20, 20))
np.testing.assert_allclose(a.tensor, a_final)
def test_fft():
n = 3
initial_state = [complex(0)] * (1 << n)
initial_state[1] = 1j
initial_state[5] = -1
initial_node = tn.Node(np.array(initial_state).reshape((2, ) * n))
fft_out = fft.add_fft([initial_node[k] for k in range(n)])
result = tn.contractors.greedy(tn.reachable(fft_out[0].node1), fft_out)
tn.flatten_edges(fft_out)
actual = result.tensor
expected = np.fft.fft(initial_state, norm="ortho")
np.testing.assert_allclose(expected, actual)
def test_flatten_trace_consistent_tensor(dtype, num_charges):
a_val = get_random_symmetric((5, 5, 5, 5, 5),
[False, False, True, True, True],
num_charges,
dtype=dtype)
a = tn.Node(a_val, backend='symmetric')
e1 = tn.connect(a[0], a[4])
e2 = tn.connect(a[3], a[2])
tn.flatten_edges([e2, e1])
tn.check_correct({a})
# Check expected values.
a_final = np.reshape(np.transpose(a_val.todense(), (1, 2, 0, 3, 4)),
(5, 25, 25))
np.testing.assert_allclose(a.tensor.todense(), a_final)
def flatten_bipartite_densitymatrix(bipartite_tensor):
""" Flatten a bipartite state density matrix to a rank-2 matrix.
:param bipartite_tensor: density matrix (rank-4) for the bipartite system
:return: flatten rank-2 density matrix
:type bipartite_tensor: numpy.ndarray
:rtype: numpy.ndarray
"""
denmat_node = tn.Node(bipartite_tensor)
e0, e1, e2, e3 = denmat_node[0], denmat_node[1], denmat_node[
2], denmat_node[3]
tn.flatten_edges([e0, e1])
tn.flatten_edges([e2, e3])
return denmat_node.tensor
def test_flatten_consistent_tensor(backend):
a_val = np.ones((2, 3, 4, 5))
b_val = np.ones((3, 5, 4, 2))
a = tn.Node(a_val, backend=backend)
b = tn.Node(b_val, backend=backend)
e1 = tn.connect(a[0], b[3])
e2 = tn.connect(b[1], a[3])
e3 = tn.connect(a[1], b[0])
tn.flatten_edges([e3, e1, e2])
tn.check_correct({a, b})
# Check expected values.
a_final = np.reshape(np.transpose(a_val, (2, 1, 0, 3)), (4, 30))
b_final = np.reshape(np.transpose(b_val, (2, 0, 3, 1)), (4, 30))
np.testing.assert_allclose(a.tensor, a_final)
np.testing.assert_allclose(b.tensor, b_final)
def test_flatten_consistent_result(backend):
a_val = np.ones((3, 5, 5, 6))
b_val = np.ones((5, 6, 4, 5))
# Create non flattened example to compare against.
a_noflat = tn.Node(a_val, backend=backend)
b_noflat = tn.Node(b_val, backend=backend)
e1 = tn.connect(a_noflat[1], b_noflat[3])
e2 = tn.connect(a_noflat[3], b_noflat[1])
e3 = tn.connect(a_noflat[2], b_noflat[0])
a_dangling_noflat = a_noflat[0]
b_dangling_noflat = b_noflat[2]
for edge in [e1, e2, e3]:
noflat_result_node = tn.contract(edge)
noflat_result_node.reorder_edges([a_dangling_noflat, b_dangling_noflat])
noflat_result = noflat_result_node.tensor
# Create network with flattening
a_flat = tn.Node(a_val, backend=backend)
b_flat = tn.Node(b_val, backend=backend)
e1 = tn.connect(a_flat[1], b_flat[3])
e2 = tn.connect(a_flat[3], b_flat[1])
e3 = tn.connect(a_flat[2], b_flat[0])
a_dangling_flat = a_flat[0]
b_dangling_flat = b_flat[2]
final_edge = tn.flatten_edges([e1, e2, e3])
flat_result_node = tn.contract(final_edge)
flat_result_node.reorder_edges([a_dangling_flat, b_dangling_flat])
flat_result = flat_result_node.tensor
np.testing.assert_allclose(flat_result, noflat_result)
def test_flatten_edges_dangling(backend):
a = tn.Node(np.zeros((2, 3, 4, 5)), name="A", backend=backend)
e1 = a[0]
e2 = a[1]
e3 = a[2]
e4 = a[3]
flattened_edge = tn.flatten_edges([e1, e3], new_edge_name="New Edge")
assert a.shape == (3, 5, 8)
assert a.edges == [e2, e4, flattened_edge]
assert flattened_edge.name == "New Edge"
tn.check_correct({a})
def test_flatten_trace_edges(backend):
a = tn.Node(np.zeros((2, 3, 4, 3, 5, 5)), backend=backend)
c = tn.Node(np.zeros((2, 4)), backend=backend)
e1 = tn.connect(a[1], a[3])
e2 = tn.connect(a[4], a[5])
external_1 = tn.connect(a[0], c[0])
external_2 = tn.connect(c[1], a[2])
new_edge = tn.flatten_edges([e1, e2], "New Edge")
tn.check_correct({a, c})
assert a.shape == (2, 4, 15, 15)
assert a.edges == [external_1, external_2, new_edge, new_edge]
assert new_edge.name == "New Edge"
def test_flatten_edges_standard(backend):
a = tn.Node(np.zeros((2, 3, 5)), name="A", backend=backend)
b = tn.Node(np.zeros((2, 3, 4, 5)), name="B", backend=backend)
e1 = tn.connect(a[0], b[0], "Edge_1_1")
e2 = tn.connect(a[2], b[3], "Edge_2_3")
edge_a_1 = a[1]
edge_b_1 = b[1]
edge_b_2 = b[2]
new_edge = tn.flatten_edges([e1, e2], new_edge_name="New Edge")
assert a.shape == (3, 10)
assert b.shape == (3, 4, 10)
assert a.edges == [edge_a_1, new_edge]
assert b.edges == [edge_b_1, edge_b_2, new_edge]
tn.check_correct({a, b})
def test_flatten_trace_consistent_result(backend):
a_val = np.ones((5, 6, 6, 7, 5, 7))
a_noflat = tn.Node(a_val, backend=backend)
e1 = tn.connect(a_noflat[0], a_noflat[4])
e2 = tn.connect(a_noflat[1], a_noflat[2])
e3 = tn.connect(a_noflat[3], a_noflat[5])
for edge in [e1, e2, e3]:
noflat_result = tn.contract(edge).tensor
# Create network with flattening
a_flat = tn.Node(a_val)
e1 = tn.connect(a_flat[0], a_flat[4])
e2 = tn.connect(a_flat[1], a_flat[2])
e3 = tn.connect(a_flat[3], a_flat[5])
final_edge = tn.flatten_edges([e1, e2, e3])
flat_result = tn.contract(final_edge).tensor
np.testing.assert_allclose(flat_result, noflat_result)
def test_flatten_trace_consistent_result(dtype, num_charges):
a_val = get_random_symmetric((5, 5, 5, 5, 5, 5),
[False, False, True, True, True, False],
num_charges,
dtype=dtype)
a_noflat = tn.Node(a_val, backend='symmetric')
e1 = tn.connect(a_noflat[0], a_noflat[4])
e2 = tn.connect(a_noflat[1], a_noflat[2])
e3 = tn.connect(a_noflat[3], a_noflat[5])
for edge in [e1, e2, e3]:
noflat_result = tn.contract(edge).tensor
# Create network with flattening
a_flat = tn.Node(a_val, backend='symmetric')
e1 = tn.connect(a_flat[0], a_flat[4])
e2 = tn.connect(a_flat[1], a_flat[2])
e3 = tn.connect(a_flat[3], a_flat[5])
final_edge = tn.flatten_edges([e1, e2, e3])
flat_result = tn.contract(final_edge).tensor
flat_result = flat_result.contiguous()
noflat_result = noflat_result.contiguous()
np.testing.assert_allclose(flat_result.data, noflat_result.data)
def test_flatten_consistent_result(dtype, num_charges):
a_val = get_random_symmetric((10, 10, 10, 10), [False] * 4,
num_charges,
dtype=dtype)
b_val = get_random_symmetric((10, 10, 10, 10), [True] * 4,
num_charges,
dtype=dtype)
# Create non flattened example to compare against.
a_noflat = tn.Node(a_val, backend='symmetric')
b_noflat = tn.Node(b_val, backend='symmetric')
e1 = tn.connect(a_noflat[1], b_noflat[3])
e2 = tn.connect(a_noflat[3], b_noflat[1])
e3 = tn.connect(a_noflat[2], b_noflat[0])
a_dangling_noflat = a_noflat[0]
b_dangling_noflat = b_noflat[2]
for edge in [e1, e2, e3]:
noflat_result_node = tn.contract(edge)
noflat_result_node.reorder_edges([a_dangling_noflat, b_dangling_noflat])
noflat_result = noflat_result_node.tensor
# Create network with flattening
a_flat = tn.Node(a_val, backend='symmetric')
b_flat = tn.Node(b_val, backend='symmetric')
e1 = tn.connect(a_flat[1], b_flat[3])
e2 = tn.connect(a_flat[3], b_flat[1])
e3 = tn.connect(a_flat[2], b_flat[0])
a_dangling_flat = a_flat[0]
b_dangling_flat = b_flat[2]
final_edge = tn.flatten_edges([e1, e2, e3])
flat_result_node = tn.contract(final_edge)
flat_result_node.reorder_edges([a_dangling_flat, b_dangling_flat])
flat_result = flat_result_node.tensor
flat_result = flat_result.contiguous()
noflat_result = noflat_result.contiguous()
np.testing.assert_allclose(flat_result.data, noflat_result.data)
def test_flatten_edges_empty_list_value_error(backend):
a = tn.Node(np.eye(2), backend=backend)
b = tn.Node(np.eye(2), backend=backend)
tn.connect(a[0], b[0])
with pytest.raises(ValueError):
tn.flatten_edges([])
def test_flatten_edges_different_backend_raises_value_error(single_node_edge):
node1 = single_node_edge.node
node2 = tn.Node(np.random.rand(2, 2, 2))
node2.backend = BaseBackend()
with pytest.raises(ValueError):
tn.flatten_edges(node1.get_all_edges() + node2.get_all_edges())