# 10. Tensor Network Training of Quantum Circuits¶

Here we’ll run through constructing a tensor network of an ansatz quantum circuit, then training certain ‘parametrizable’ tensors representing quantum gates in that tensor network to replicate the behaviour of a target unitary.

%config InlineBackend.figure_formats = ['svg']
import quimb as qu
import quimb.tensor as qtn


## 10.1. The Ansatz Circuit¶

First we set up the ansatz circuit and extract the tensor network. Key here is that when we supply parametrize=True to the 'U3' gate call, it injects a PTensor into the network, which lazily represents its data array with a function and set of parameters. Later, when the optimizer sees this it then knows to optimize the parameters rather than the array itself.

def single_qubit_layer(circ, gate_round=None):
"""Apply a parametrizable layer of single qubit U3 gates.
"""
for i in range(circ.N):
# initialize with random parameters
params = qu.randn(3, dist='uniform')
circ.apply_gate(
'U3', *params, i,
gate_round=gate_round, parametrize=True)

def two_qubit_layer(circ, gate2='CZ', reverse=False, gate_round=None):
"""Apply a layer of constant entangling gates.
"""
regs = range(0, circ.N - 1)
if reverse:
regs = reversed(regs)

for i in regs:
circ.apply_gate(
gate2, i, i + 1, gate_round=gate_round)

def ansatz_circuit(n, depth, gate2='CZ', **kwargs):
"""Construct a circuit of single qubit and entangling layers.
"""
circ = qtn.Circuit(n, **kwargs)

for r in range(depth):
# single qubit gate layer
single_qubit_layer(circ, gate_round=r)

# alternate between forward and backward CZ layers
two_qubit_layer(
circ, gate2=gate2, gate_round=r, reverse=r % 2 == 0)

# add a final single qubit layer
single_qubit_layer(circ, gate_round=r + 1)

return circ


The form of the 'U3' gate (which generalizes all possible single qubit gates) can be seen here - U_gate(). Now we are ready to instantiate a circuit:

n = 6
depth = 9
gate2 = 'CZ'

circ = ansatz_circuit(n, depth, gate2=gate2)
circ

<Circuit(n=6, num_gates=105, gate_opts={'contract': 'auto-split-gate', 'propagate_tags': 'register'})>


We can extract just the unitary part of the circuit as a tensor network like so:

V = circ.uni

/media/johnnie/Storage2TB/Sync/dev/python/quimb/quimb/tensor/circuit.py:1194: FutureWarning: In future the tensor network returned by circ.uni will not be transposed as it is currently, to match the expectation from U = circ.uni.to_dense() behaving like U @ psi. You can retain this behaviour with circ.get_uni(transposed=True).
warnings.warn(


You can see it already has various tags simultaneously identifying different structures:

# types of gate
V.draw(color=['U3', gate2], show_inds=True)

# layers of gates
V.draw(color=[f'ROUND_{i}' for i in range(depth + 1)], show_inds=True)

# what register each tensor is 'above'
V.draw(color=[f'I{i}' for i in range(n)], show_inds=True)

# a unique tag for per gate applied
V.draw(color=[f'GATE_{i}' for i in range(circ.num_gates)], legend=False)


## 10.2. The Target Unitary¶

Next we need a target unitary to try and digitially replicate. Here we’ll take an Ising Hamiltonian and a short time evolution. Once we have the dense (matrix) form of the target unitary $U$ we need to convert it to a tensor which we can put in a tensor network:

# the hamiltonian
H = qu.ham_ising(n, jz=1.0, bx=0.7, cyclic=False)

# the propagator for the hamiltonian
t = 2
U_dense = qu.expm(-1j * t * H)

# 'tensorized' version of the unitary propagator
U = qtn.Tensor(
data=U_dense.reshape([2] * (2 * n)),
inds=[f'k{i}' for i in range(n)] + [f'b{i}' for i in range(n)],
tags={'U_TARGET'}
)
U.draw(color=['U3', gate2, 'U_TARGET'])


The core object describing how similar two unitaries are is: $$\mathrm{Tr}(V^{\dagger}U)$$, which we can naturally visualize at a tensor network:

(V.H & U).draw(color=['U3', gate2, 'U_TARGET'])


For our loss function we’ll normalize this and negate it (since the optimizer minimizes).

def loss(V, U):
return 1 - abs((V.H & U).contract(all, optimize='auto-hq')) / 2**n

# check our current unitary 'infidelity':
loss(V, U)

0.9881326237593226


So as expected currently the two unitaries are not similar at all.

## 10.3. The Tensor Network Optimization¶

Now we are ready to construct the TNOptimizer object, with options detailed below:

# use the autograd/jax based optimizer

tnopt = qtn.TNOptimizer(
V,                        # the tensor network we want to optimize
loss,                     # the function we want to minimize
loss_constants={'U': U},  # supply U to the loss function as a constant TN
tags=['U3'],              # only optimize U3 tensors
autodiff_backend='jax',   # use 'autograd' for non-compiled optimization
optimizer='L-BFGS-B',     # the optimization algorithm
)


Note

If tags is not specified the default is to optimize all tensors. In that case, instead of specifying tensor tags to opt-in you can use constant_tags to opt-out tensors tags you don’t want to optimize, which may be more convenient.

We could call optimize for pure gradient based optimization, but since unitary circuits can be tricky we’ll use optimize_basinhopping which combines gradient descent with ‘hopping’ to escape local minima:

# allow 10 hops with 500 steps in each 'basin'
V_opt = tnopt.optimize_basinhopping(n=500, nhop=10)

+0.005301594734 [best: +0.005261898041] :  45%|████▌     | 2258/5000 [00:40<00:49, 55.26it/s]


The optimized tensor network still contains PTensor instances but now with optimized parameters. For example, here’s the tensor of the U3 gate acting on qubit-2 in round-4:

V_opt['U3', 'I2', 'ROUND_4']

PTensor(shape=(2, 2), inds=('_8d3684AAABr', '_8d3684AAABh'), tags=oset(['GATE_46', 'ROUND_4', 'U3', 'I2']))


We can see the parameters have been updated by the training:

# the initial values
V['U3', 'ROUND_4', 'I2'].params

array([0.25927184, 0.10598236, 0.42593174])

# the optimized values
V_opt['U3', 'ROUND_4', 'I2'].params

array([ 0.52709323,  0.6036498 , -0.23340225], dtype=float32)


We can see what gate these parameters would generate:

qu.U_gate(*V_opt['U3', 'ROUND_4', 'I2'].params)

[[ 0.965472-0.j       -0.253443+0.060252j]
[ 0.214467+0.147877j  0.90005 +0.349352j]]


A final sanity check we can perform is to try evolving a random state with the target unitary and trained circuit and check the fidelity between the resulting states.

First we turn the tensor network version of $$V$$ into a dense matrix:

V_opt_dense = V_opt.to_dense([f'k{i}' for i in range(n)], [f'b{i}' for i in range(n)])


Next we create a random initial state, and evolve it with the

psi0 = qu.rand_ket(2**n)

# this is the exact state we want
psif_exact = U_dense @ psi0

# this is the state our circuit will produce if fed psi0
psif_apprx = V_opt_dense @ psi0


The (in)fidelity should broadly match our training loss:

f"Fidelity: {100 * qu.fidelity(psif_apprx, psif_exact):.2f} %"

'Fidelity: 99.57 %'


## 10.4. Extracting the New Circuit¶

We can extract the trained circuit parameters by updating the original Circuit object from the trained TN:

circ.update_params_from(V_opt)

# the basic gate specification
circ.gates

[<Gate(label=U3, params=(0.16929962, -0.6534094, -0.25308686), qubits=(0,), round=0, parametrize=True))>,
<Gate(label=U3, params=(-0.8037065, -0.25300127, 1.5709982), qubits=(1,), round=0, parametrize=True))>,
<Gate(label=U3, params=(-1.0775858, 0.29377267, 1.573218), qubits=(2,), round=0, parametrize=True))>,
<Gate(label=U3, params=(-0.70041966, -0.54078287, 1.5730082), qubits=(3,), round=0, parametrize=True))>,
<Gate(label=U3, params=(2.233548, -0.6997934, 1.5887524), qubits=(4,), round=0, parametrize=True))>,
<Gate(label=U3, params=(2.6693592, -0.3773012, 1.6181816), qubits=(5,), round=0, parametrize=True))>,
<Gate(label=CZ, params=(), qubits=(4, 5), round=0)>,
<Gate(label=CZ, params=(), qubits=(3, 4), round=0)>,
<Gate(label=CZ, params=(), qubits=(2, 3), round=0)>,
<Gate(label=CZ, params=(), qubits=(1, 2), round=0)>,
<Gate(label=CZ, params=(), qubits=(0, 1), round=0)>,
<Gate(label=U3, params=(0.9794985, -0.1482388, 0.3628862), qubits=(0,), round=1, parametrize=True))>,
<Gate(label=U3, params=(1.3343341, -0.8915439, 0.252669), qubits=(1,), round=1, parametrize=True))>,
<Gate(label=U3, params=(0.0016000561, 1.3208843, 1.8358507), qubits=(2,), round=1, parametrize=True))>,
<Gate(label=U3, params=(0.20370977, -0.5142169, 0.52689326), qubits=(3,), round=1, parametrize=True))>,
<Gate(label=U3, params=(1.7909783, 0.082318805, 0.6477436), qubits=(4,), round=1, parametrize=True))>,
<Gate(label=U3, params=(-1.6923836, 0.22169495, -0.06101191), qubits=(5,), round=1, parametrize=True))>,
<Gate(label=CZ, params=(), qubits=(0, 1), round=1)>,
<Gate(label=CZ, params=(), qubits=(1, 2), round=1)>,
<Gate(label=CZ, params=(), qubits=(2, 3), round=1)>,
<Gate(label=CZ, params=(), qubits=(3, 4), round=1)>,
<Gate(label=CZ, params=(), qubits=(4, 5), round=1)>,
<Gate(label=U3, params=(-1.5676986, 0.79075044, 1.9245409), qubits=(0,), round=2, parametrize=True))>,
<Gate(label=U3, params=(1.4227079, 0.50024813, 0.8913113), qubits=(1,), round=2, parametrize=True))>,
<Gate(label=U3, params=(-1.5313544, 0.5816354, -0.30510747), qubits=(2,), round=2, parametrize=True))>,
<Gate(label=U3, params=(0.05509888, 0.46960974, 0.5104008), qubits=(3,), round=2, parametrize=True))>,
<Gate(label=U3, params=(-1.4967402, 2.4865503, -0.09552598), qubits=(4,), round=2, parametrize=True))>,
<Gate(label=U3, params=(1.6648741, 1.5529978, 0.33650705), qubits=(5,), round=2, parametrize=True))>,
<Gate(label=CZ, params=(), qubits=(4, 5), round=2)>,
<Gate(label=CZ, params=(), qubits=(3, 4), round=2)>,
<Gate(label=CZ, params=(), qubits=(2, 3), round=2)>,
<Gate(label=CZ, params=(), qubits=(1, 2), round=2)>,
<Gate(label=CZ, params=(), qubits=(0, 1), round=2)>,
<Gate(label=U3, params=(2.3430626, 0.5120296, 0.76947355), qubits=(0,), round=3, parametrize=True))>,
<Gate(label=U3, params=(1.4037384, 0.34615904, 2.6413994), qubits=(1,), round=3, parametrize=True))>,
<Gate(label=U3, params=(-1.4830327, 0.22897828, -0.58199924), qubits=(2,), round=3, parametrize=True))>,
<Gate(label=U3, params=(0.002939354, 0.941026, 2.3286314), qubits=(3,), round=3, parametrize=True))>,
<Gate(label=U3, params=(-0.19347823, 1.4644276, 0.6541145), qubits=(4,), round=3, parametrize=True))>,
<Gate(label=U3, params=(0.79052246, 1.3951869, -0.1806278), qubits=(5,), round=3, parametrize=True))>,
<Gate(label=CZ, params=(), qubits=(0, 1), round=3)>,
<Gate(label=CZ, params=(), qubits=(1, 2), round=3)>,
<Gate(label=CZ, params=(), qubits=(2, 3), round=3)>,
<Gate(label=CZ, params=(), qubits=(3, 4), round=3)>,
<Gate(label=CZ, params=(), qubits=(4, 5), round=3)>,
<Gate(label=U3, params=(1.7128671, -0.83896357, 0.87925446), qubits=(0,), round=4, parametrize=True))>,
<Gate(label=U3, params=(1.7145393, -0.38567403, -0.34633058), qubits=(1,), round=4, parametrize=True))>,
<Gate(label=U3, params=(0.52709323, 0.6036498, -0.23340225), qubits=(2,), round=4, parametrize=True))>,
<Gate(label=U3, params=(0.005033329, -0.7817636, -0.5969079), qubits=(3,), round=4, parametrize=True))>,
<Gate(label=U3, params=(2.1565654, 0.751235, -1.485796), qubits=(4,), round=4, parametrize=True))>,
<Gate(label=U3, params=(1.5731558, 1.591097, 0.31102985), qubits=(5,), round=4, parametrize=True))>,
<Gate(label=CZ, params=(), qubits=(4, 5), round=4)>,
<Gate(label=CZ, params=(), qubits=(3, 4), round=4)>,
<Gate(label=CZ, params=(), qubits=(2, 3), round=4)>,
<Gate(label=CZ, params=(), qubits=(1, 2), round=4)>,
<Gate(label=CZ, params=(), qubits=(0, 1), round=4)>,
<Gate(label=U3, params=(-2.325529, 0.418274, -1.1409844), qubits=(0,), round=5, parametrize=True))>,
<Gate(label=U3, params=(0.011331396, -0.6176183, 0.28041682), qubits=(1,), round=5, parametrize=True))>,
<Gate(label=U3, params=(1.4015418, 0.5063352, -0.5949921), qubits=(2,), round=5, parametrize=True))>,
<Gate(label=U3, params=(0.009612298, 0.2217627, 0.04534868), qubits=(3,), round=5, parametrize=True))>,
<Gate(label=U3, params=(1.8638043, -0.20229527, -0.7953688), qubits=(4,), round=5, parametrize=True))>,
<Gate(label=U3, params=(-0.02179846, -0.57399744, 0.8407001), qubits=(5,), round=5, parametrize=True))>,
<Gate(label=CZ, params=(), qubits=(0, 1), round=5)>,
<Gate(label=CZ, params=(), qubits=(1, 2), round=5)>,
<Gate(label=CZ, params=(), qubits=(2, 3), round=5)>,
<Gate(label=CZ, params=(), qubits=(3, 4), round=5)>,
<Gate(label=CZ, params=(), qubits=(4, 5), round=5)>,
<Gate(label=U3, params=(1.3877774, 0.9551876, 0.9025401), qubits=(0,), round=6, parametrize=True))>,
<Gate(label=U3, params=(0.002978851, 2.6452973, 0.20050335), qubits=(1,), round=6, parametrize=True))>,
<Gate(label=U3, params=(0.44300193, -0.30685717, -0.5065869), qubits=(2,), round=6, parametrize=True))>,
<Gate(label=U3, params=(0.076909065, 1.6145856, 0.6035761), qubits=(3,), round=6, parametrize=True))>,
<Gate(label=U3, params=(-0.5196627, 1.6240075, 0.20132811), qubits=(4,), round=6, parametrize=True))>,
<Gate(label=U3, params=(2.2904618, 0.9574056, -0.4791421), qubits=(5,), round=6, parametrize=True))>,
<Gate(label=CZ, params=(), qubits=(4, 5), round=6)>,
<Gate(label=CZ, params=(), qubits=(3, 4), round=6)>,
<Gate(label=CZ, params=(), qubits=(2, 3), round=6)>,
<Gate(label=CZ, params=(), qubits=(1, 2), round=6)>,
<Gate(label=CZ, params=(), qubits=(0, 1), round=6)>,
<Gate(label=U3, params=(-0.30187988, -0.87076926, -0.8777541), qubits=(0,), round=7, parametrize=True))>,
<Gate(label=U3, params=(1.624832, 1.3061262, 1.019058), qubits=(1,), round=7, parametrize=True))>,
<Gate(label=U3, params=(-0.0049230885, 1.0577729, 1.851811), qubits=(2,), round=7, parametrize=True))>,
<Gate(label=U3, params=(0.02848778, 1.4706283, 1.2345053), qubits=(3,), round=7, parametrize=True))>,
<Gate(label=U3, params=(0.2926683, -0.228561, 0.6961813), qubits=(4,), round=7, parametrize=True))>,
<Gate(label=U3, params=(-0.64344996, -0.6371593, -1.0523038), qubits=(5,), round=7, parametrize=True))>,
<Gate(label=CZ, params=(), qubits=(0, 1), round=7)>,
<Gate(label=CZ, params=(), qubits=(1, 2), round=7)>,
<Gate(label=CZ, params=(), qubits=(2, 3), round=7)>,
<Gate(label=CZ, params=(), qubits=(3, 4), round=7)>,
<Gate(label=CZ, params=(), qubits=(4, 5), round=7)>,
<Gate(label=U3, params=(-1.3433273, 1.4197886, 0.5518878), qubits=(0,), round=8, parametrize=True))>,
<Gate(label=U3, params=(-1.7108984, -0.03164328, -1.3058407), qubits=(1,), round=8, parametrize=True))>,
<Gate(label=U3, params=(-1.4684929, -0.23041597, 0.5393066), qubits=(2,), round=8, parametrize=True))>,
<Gate(label=U3, params=(0.008398809, -0.18387741, -0.6591129), qubits=(3,), round=8, parametrize=True))>,
<Gate(label=U3, params=(-0.000680109, 0.2461698, -0.4091857), qubits=(4,), round=8, parametrize=True))>,
<Gate(label=U3, params=(0.110012054, 1.3755561, 1.2930975), qubits=(5,), round=8, parametrize=True))>,
<Gate(label=CZ, params=(), qubits=(4, 5), round=8)>,
<Gate(label=CZ, params=(), qubits=(3, 4), round=8)>,
<Gate(label=CZ, params=(), qubits=(2, 3), round=8)>,
<Gate(label=CZ, params=(), qubits=(1, 2), round=8)>,
<Gate(label=CZ, params=(), qubits=(0, 1), round=8)>,
<Gate(label=U3, params=(0.4610416, 0.26438755, 1.6497571), qubits=(0,), round=9, parametrize=True))>,
<Gate(label=U3, params=(-0.7864091, 1.571062, 0.03173134), qubits=(1,), round=9, parametrize=True))>,
<Gate(label=U3, params=(-1.0439473, 1.5646937, 0.2356552), qubits=(2,), round=9, parametrize=True))>,
<Gate(label=U3, params=(0.88693166, 1.5739466, -0.4111447), qubits=(3,), round=9, parametrize=True))>,
<Gate(label=U3, params=(0.4123549, 1.0429444, 1.6132663), qubits=(4,), round=9, parametrize=True))>,
<Gate(label=U3, params=(1.3669215, 3.024821, 0.9148983), qubits=(5,), round=9, parametrize=True))>]