Tensor networks

A tensor network is a graph where each node represents a tensor and each edge represents a shared index between tensors. In the Tensor section, it was shown that a tensor can be graphically modelled as a circle with a leg for each index of the tensor. It follows that a tensor network can be represented as a collection of circles and lines where stray legs denote free indices:

One of the key operations that can be performed over a tensor network is a contraction. Local contractions are discussed in the Tensor section and combine two tensors which share at least one index in a tensor network. Global contractions reduce a tensor network to a single node by iteratively performing local contractions until all of the edges in the tensor network are consumed. For example:

Above, the result of the tensor network contraction is the node \(E\). Observe that \(E\) was produced by first contracting nodes \(C\) and \(D\) to create node \(CD\), then contracting node \(CD\) with node \(B\) to generate node \(BCD\), and then finally contracting node \(BCD\) with node \(A\). Here, nodes \(CD\) and \(BCD\) are intermediary tensors and the order of contractions is summarized by the contraction path

\[P = \{(C, D), (CD, B), (BCD, A)\}\,.\]

In general, the contraction path of a tensor network is not unique and has a significant impact on the memory requirements and running time of a tensor network contraction.

Modelling quantum circuits

Although tensor networks are pure mathematical objects, they are often used to model and simulate quantum circuits. For example, consider the following circuit which generates an EPR pair from two unentangled \(\vert 0 \rangle\) qubits:

This circuit can be directly modelled with the following tensor network:

To construct this tensor network in Jet, it is necessary to first define each of the consituent tensors using the Tensor class. Recall that:

\[\begin{split}\vert 0 \rangle = \begin{bmatrix} 1 \\ 0 \end{bmatrix} \qquad H = \begin{bmatrix} \frac{1}{\sqrt{2}} & \frac{1}{\sqrt{2}} \\ \frac{1}{\sqrt{2}} & -\frac{1}{\sqrt{2}} \end{bmatrix} \qquad CNOT = \begin{bmatrix} 1 & 0 & 0 & 0 \\ 0 & 1 & 0 & 0 \\ 0 & 0 & 0 & 1 \\ 0 & 0 & 1 & 0 \end{bmatrix} \,.\end{split}\]

The \(\vert 0 \rangle\) qubits are relatively simple to create:

C++Python

using Tensor = Jet::Tensor<std::complex<float>>;

// The control qubit is defined as a 1-D vector with 2 elements.
Tensor q0({"i"}, {2}, {1, 0});

// Note that the index of ``q1`` differs from ``q0``.
Tensor q1({"j"}, {2}, {1, 0});

The Hadamard gate \(H\) can also be constructed in the usual way:

C++Python

const float inv_sqrt_2 = 1 / std::sqrt(2);
Tensor H({"i", "k"}, {2, 2}, {inv_sqrt_2, inv_sqrt_2, inv_sqrt_2, -inv_sqrt_2});

The controlled NOT gate \(CNOT\) is slightly trickier. From the diagram, \(CNOT \in \mathbb{C}^{2 \times 2 \times 2 \times 2}\). To derive this \(CNOT\) tensor, note that a two-qubit state \(\vert \psi \rangle \in \mathbb{C}^{4}\) can be encoded as a \(\mathbb{C}^{2 \times 2}\) matrix:

\[\begin{split}\vert \psi \rangle = \alpha_{00} \vert 00 \rangle + \alpha_{01} \vert 01 \rangle + \alpha_{10} \vert 10 \rangle + \alpha_{11} \vert 11 \rangle = \begin{bmatrix} \alpha_{00} & \alpha_{01} \\ \alpha_{10} & \alpha_{11} \end{bmatrix}\,.\end{split}\]

It follows that

\[\begin{split}CNOT_{0, 0} = \begin{bmatrix} 1 & 0 \\ 0 & 0 \end{bmatrix} \quad CNOT_{0, 1} = \begin{bmatrix} 0 & 1 \\ 0 & 0 \end{bmatrix} \quad CNOT_{1, 0} = \begin{bmatrix} 0 & 0 \\ 0 & 1 \end{bmatrix} \quad CNOT_{1, 1} = \begin{bmatrix} 0 & 0 \\ 1 & 0 \end{bmatrix}\,.\end{split}\]

The \(CNOT\) gate is then given by

C++Python

Tensor CNOT({"k", "j", "m", "n"}, {2, 2, 2, 2});
CNOT.SetValue({0, 0, 0, 0}, 1); // |00> -> |00>
CNOT.SetValue({0, 1, 0, 1}, 1); // |01> -> |01>
CNOT.SetValue({1, 0, 1, 1}, 1); // |10> -> |11>
CNOT.SetValue({1, 1, 1, 0}, 1); // |11> -> |10>

Now, creating the tensor network is easy with the TensorNetwork class:

C++Python

using TensorNetwork = Jet::TensorNetwork<Tensor>;
TensorNetwork tn;

// The second argument can be used to associate a tensor with a set of tags.
tn.AddTensor(q0);
tn.AddTensor(q1);
tn.AddTensor(H);
tn.AddTensor(CNOT);

By default, the TensorNetwork class performs contractions in random order:

C++Python

tn.Contract();

An explicit contraction path can also be specified by providing a list of pair of node IDs (0-indexed) to the Contract() function. The ID of a node is the order in which it was added to the tensor network. Intermediate tensors are assigned node IDs according to SSA convention (i.e., they are assigned the node ID immediately following the largest node ID in the tensor network in use at the time the intermediate tensor was created).

C++Python

tn.Contract({{0, 2}, {1, 3}, {4, 5}});

Putting it all together,

C++PythonPython (using the XIR)

#include <cmath>
#include <complex>
#include <iostream>

#include <Jet.hpp>

int main()
{
    using Tensor = Jet::Tensor<std::complex<float>>;
    using TensorNetwork = Jet::TensorNetwork<Tensor>;

    Tensor q0({"i"}, {2}, {1, 0});
    Tensor q1({"j"}, {2}, {1, 0});

    const float inv_sqrt_2 = 1 / std::sqrt(2);
    Tensor H({"i", "k"}, {2, 2}, {inv_sqrt_2, inv_sqrt_2, inv_sqrt_2, -inv_sqrt_2});

    Tensor CNOT({"k", "j", "m", "n"}, {2, 2, 2, 2});
    CNOT.SetValue({0, 0, 0, 0}, 1);
    CNOT.SetValue({0, 1, 0, 1}, 1);
    CNOT.SetValue({1, 0, 1, 1}, 1);
    CNOT.SetValue({1, 1, 1, 0}, 1);

    TensorNetwork tn;
    tn.AddTensor(q0);
    tn.AddTensor(q1);
    tn.AddTensor(H);
    tn.AddTensor(CNOT);

    Tensor result = tn.Contract();
    std::cout << "Amplitude |00> = " << result.GetValue({0, 0}) << std::endl;
    std::cout << "Amplitude |01> = " << result.GetValue({0, 1}) << std::endl;
    std::cout << "Amplitude |10> = " << result.GetValue({1, 0}) << std::endl;
    std::cout << "Amplitude |11> = " << result.GetValue({1, 1}) << std::endl;

    return 0;
}

The output of the program is

C++Python

Amplitude |00> = (0.707107,0)
Amplitude |01> = (0,0)
Amplitude |10> = (0,0)
Amplitude |11> = (0.707107,0)

Task-based contraction

While TensorNetwork::Contract() is simple to use, it is unlikely to exhibit optimal performance for large tensor networks. One alternative to the vanilla tensor network contractor is the TaskBasedContractor class which models a tensor network contraction as a parallel task scheduling problem where each task encapsulates a local tensor contraction. Such a formulation enables intermediate tensors which do not depend on each another to be contracted concurrently. As an example, consider the task graph for the quantum circuit described in the previous section:

Clearly, the leftmost nodes in the top row (\(\vert 0 \rangle\) and \(CNOT\)) may be contracted in parallel with the rightmost nodes in the top row (the other \(\vert 0 \rangle\) and \(H\)); however, the contraction representing the final output of the circuit may only be performed once nodes \(A_k\) and \(B_{m,n,k}\) have been computed.

Despite its underlying complexity, the interface to TaskBasedContractor is relatively straightforward. After constructing the TensorNetwork in the previous section, the contraction path is specified using a PathInfo object:

C++Python

PathInfo path_info(tn, {{0, 2}, {1, 3}, {4, 5}});

The contraction tasks can then be added to a new TaskBasedContractor instance:

C++Python

TaskBasedContractor<Tensor<std::complex<float>>> tbc;
tbc.AddContractionTasks(tn, path_info);

Finally, TaskBasedContractor::Contract() launches the contraction and returns a future that becomes available when the contraction is complete:

C++Python

// Start the tensor network contraction and wait for it to finish.
auto future = tbc.Contract();
future.wait();

// Each call to AddContractionTasks() generates a new result.
const auto results = tbc.GetResults();
const auto result = results[0];

use/tensor_networks

 

Download Python script

 

Download Notebook

 

View on GitHub