Variational quantum circuit with repeated unitaries

[ghost]

I would like to variationally learn a diagonalization of a short time Hamiltonian e.g. exactly what is being done in your example 10 'Tensor Network Training of Quantum Circuits', but with an ansatz of the form WDW^{\dagger} (see eqn (1) of https://arxiv.org/pdf/1910.04292.pdf).
In Quimb you do not explicitly give parameters as an input to the parameterised unitary, so I am unsure how I would 'copy' the W unitary and apply its inverse later.
Another tangential reason this would be useful would be training a translationally invariant quantum circuit, where you are training one cell per layer that is repeated multiple times.

[jcmgray]

Hi @JoeGibbs88, this should be possible a few ways using qtn.TNOptimizer (better docs for this including shared_tags are on their way...):

If you have the TN of the circuit W D, and target U you could either:

1. In the loss_fn, do something like Wdag = tn.select('W').conj(), and do some index rewiring to extend the circuit and form Tr [ W D Wdag U] etc.

2. Form the TN for Tr [ W D Wdag U] directly, then use shared_tags to link parameters in W and Wdag, (this specifies that tensors with the same tags share the same data or parameters). You'd just need to explicitly add the conjugation of the dagger tensors in the norm_fn e.g. before computing the contraction

In Quimb you do not explicitly give parameters as an input to the parameterised unitary

Note you can specify PTensor instances with an arbitrary function and parameters, the circuit interface just automates this process. When you call something like t.conj() it just chains the new function so in both cases above the parent parameters are still shared.

[ghost]

Hi @jcmgray, thanks for the reply. The first method sounds preferable if it means I could avoid having to generate a reversed ansatz and then linking the parameters, although I am new to Quimb so I am unfamiliar with the indexing system. Below is a naive attempt at this:

import quimb as qu
import quimb.tensor as qtn

def single_qubit_layer(circ, gate_round=None):
    
    for i in range(circ.N):
        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):
    
    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 W_circ(n, depth, gate2='CZ', **kwargs):
    
    circ = qtn.Circuit(n, **kwargs)

    for r in range(depth):
        single_qubit_layer(circ, gate_round=r)

        two_qubit_layer(
            circ, gate2=gate2, gate_round=r, reverse=r % 2 == 0)

    single_qubit_layer(circ, gate_round=r + 1)

    return circ

def D_circ(n, **kwargs):

    circ = qtn.Circuit(n, **kwargs)

    params = qu.randn(n, dist='uniform')
    for i in range(n):
        circ.apply_gate('Rz', params[i], i, parametrize=True)

    return circ

n = 2
depth = 2
gate2 = 'CZ'

W = W_circ(n, depth, gate2=gate2, tags='W')
D = D_circ(n, tags='D')

V = W.uni & D.uni & W.get_uni(transposed=True)

H = qu.ham_XY(n, jxy=1.0, bz=0.0, cyclic=False)

t = 0.1
U_dense = qu.expm(-1j * t * H)

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'}
)

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

tnopt = qtn.TNOptimizer(
    V,                        
    loss,                     
    loss_constants={'U': U},  
    tags=['U3', 'RZ'],        
    autodiff_backend='jax',   
    optimizer='L-BFGS-B',   
)

V_opt = tnopt.optimize_basinhopping(n=10**3, nhop=10, tol=10**-6)

Which returns the ValueError

ValueError: The index k0 appears more than twice! If this is intentionally a 'hyper' tensor network you will need to explicitly supply `output_inds` when contracting for example.

Do W and Wdag need to be given independent indices? I couldn't see how to do this in the source code

[jcmgray]

Yes, once you have extracted the TNs representing the unitaries from the circuit objects, the way they compose is defined by the way edges are indexed. Since they have matching outer indices, if you compose two (A & B), you get Tr[A B], but more than two makes the indices clash.

What you need to do is align them like thus:

W = W_circ(n, depth, gate2=gate2, tags='W').get_uni()
D = D_circ(n, tags='D').get_uni()
Wdag = W.conj()

# for visualizing
W.add_tag('W')
D.add_tag('D')
Wdag.add_tag('Wdag')

W.upper_ind_id = '__a{}__'
W.lower_ind_id = '__b{}__'
D.upper_ind_id = '__b{}__'
D.lower_ind_id = '__c{}__'
# transpose by switching lower and upper!
Wdag.lower_ind_id = '__c{}__'
Wdag.upper_ind_id = '__a{}__'

(W | D | Wdag).draw(['W', 'D', 'Wdag'], show_inds='all')

There is the tensor_network_align function, that I might add a trace switch too.

[ghost]

How should I combine the ansatze ready for optimization? When I sub your code snippet into the code I first posted and combine them with an & I get a TypeError from the abs() in the loss function: TypeError: bad operand type for abs(): 'Tensor'

W = W_circ(n, depth, gate2=gate2, tags='W').get_uni()
D = D_circ(n, tags='D').get_uni()
Wdag = W.conj()

# for visualizing
W.add_tag('W')
D.add_tag('D')
Wdag.add_tag('Wdag')

W.upper_ind_id = '__a{}__'
W.lower_ind_id = '__b{}__'
D.upper_ind_id = '__b{}__'
D.lower_ind_id = '__c{}__'
# transpose by switching lower and upper!
Wdag.lower_ind_id = '__c{}__'
Wdag.upper_ind_id = '__a{}__'

V = W & D & Wdag

H = qu.ham_XY(n, jxy=1.0, bz=0.0, cyclic=False)

t = 0.1
U_dense = qu.expm(-1j * t * H)

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'}
)

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

tnopt = qtn.TNOptimizer(
    V,                        
    loss,                     
    loss_constants={'U': U},  
    tags=['U3', 'RZ'],        
    autodiff_backend='jax',   
    optimizer='L-BFGS-B',   
)

[jcmgray]

What I posted constructs Tr[W D Wdag], and has no inds matching the U tensor constructed.
You just need to extend the logic above to include U, something like:

...
Wdag.lower_ind_id = '__c{}__'
Wdag.upper_ind_id = '__d{}__'
...
U = qtn.Tensor(
    data=U_dense.reshape([2] * (2 * n)),
    # explicitly put in lower and upper indices to close the trace
    inds=[f'__d{i}__' for i in range(n)] + [f'__a{i}__' for i in range(n)],
    tags={'U_TARGET'}
)
...

If you plot (W | D | Wdag | U) you should be able to check that a scalar (no dangling indices) has been formed.

[ghost]

Awesome, thanks for the help!