[ITensors] Fix `rrule` for applying quantum circuit to non-hermitian MPO with `apply_dag=true`

[ntausend]

Description

This Pull request fixes the rrule for applying a quantum circuit to an MPO which is not hermitian with the keyword apply_dag = true. Instead of assuming the MPO to be hermitian, I added an if statement in Line 122 of src/src/ITensorChainRules/mps/abstractmps.jl, checking if the MPO is hermitian. In case the MPO is hermitian, the old behavior is applied, in case the MPO is not hermitian, I perform the correct pullback, which involves two contributions, one from the normal application of the gate, and one coming from the daggered application of the gate.

Fixes #1314

For demonstration, I used a (modified) version of the code snipped provided in the Issue. In the beginning one can choose using Complex or Floats, QN conserved states or dense states. In all cases the output is the expected one.

Code snipped

using ITensors
using Zygote: pullback
using Random: seed!

seed!(1234)
N = 2
# Field type
elT = ComplexF64
# explicitly make the MPO hermitian
hermitice = false
# using QN number conservation
qn_version = false

# dagger of mpo
dagger(A::MPO) = swapprime(dag(A), 0 => 1)
# hermitian part
hermitianpart(A::MPO) = 1/2 * (A + dagger(A))
# printing only 3 digits
round_mat(M) = map(m -> round(m, digits = 3), M)

if qn_version
	qvals = [QN("Parity", 0, 2) => N÷2,
			 QN("Parity", 1, 2) => N÷2]
	idx = Index(qvals)
else
	idx = Index(N)
end

# printing only 3 digits
round_mat(M) = map(m -> round(m, digits = 3), M)

# generate the test MPO
A = MPO(randomITensor(elT, prime(idx), dag(idx)), [idx])
hermitice && (A = hermitianpart(A))

# ITensor version
Ait = contract(A)
# fix the order to compare to numerical results
inds_order = inds(Ait)
# matrix version
Am = matrix(Ait)

# apply itensor
G = randomITensor(elT, prime(idx), dag(idx))
inds_order_G = inds(G)
# matrix version
Gm = matrix(G)

# generate a dual vector MPO
M = MPO(randomITensor(elT, prime(idx), dag(idx)), [idx])
hermitice && (M = hermitianpart(M))
# ITensor version
Mit = contract(M)
# matrix version
Mm = matrix(Mit)

# apply function without dagger apply
f(G) = apply([G], A)
fit(G) = apply([G], Ait)
fm(G) = G * Am

# compute the actions
fG   = matrix(permute(contract(f(G)), inds_order))
fGit = matrix(permute(fit(G), inds_order))
fGm  = fm(Gm) 

# should all result in the same matrix
println("Results from applying gate to MPO/ITensor/Matrix")
println("Are all versions of computing the contraction results in the same result? $((fG ≈ fGit) && (fGit ≈ fGm))")

# now we calculate the pullback of all these versions onto a test dual vetor
∇fG   = matrix(permute(pullback(f, G)[2](M)[1], inds_order_G))
∇fGit = matrix(permute(pullback(fit, G)[2](Mit)[1], inds_order_G))
∇fGm  =        pullback(fm, Gm)[2](Mm)[1]

println("Is the ITensor pullback equal to the matrix pullback? $(∇fGit ≈ ∇fGm)")
println("Is the MPO pullback equal to the matrix pullback? $(∇fG ≈ ∇fGm)")


# apply function with dagger apply
g(G) = apply([G], A; apply_dag = true)
git(G) = apply([G], Ait; apply_dag = true)
gm(G) = G * Am * G'

# compute the actions
gG   = matrix(permute(contract(g(G)), inds_order))
gGit = matrix(permute(git(G), inds_order))
gGm  = gm(Gm) 

# should all result in the same matrix
println("Results from applying gate to MPO/ITensor/Matrix")
println("Are all versions of computing the contraction results in the same result? $((gG ≈ gGit) && (gGit ≈ gGm))")

# now we calculate the pullback of all these versions onto a test dual vetor
∇gG   = matrix(permute(pullback(g, G)[2](M)[1], inds_order_G))
∇gGit = matrix(permute(pullback(git, G)[2](Mit)[1], inds_order_G))
∇gGm  =        pullback(gm, Gm)[2](Mm)[1]

println("Is the ITensor pullback equal to the matrix pullback? $(∇gGit ≈ ∇gGm)")
println("Is the MPO pullback equal to the matrix pullback? $(∇gG ≈ ∇gGm)")
println(∇gGit)
println(∇gG)

Minimal demonstration of previous behavior

The output before the fix was:

Results from applying gate to MPO/ITensor/Matrix
Are all versions of computing the contraction results in the same result? true
Is the ITensor pullback equal to the matrix pullback? true
Is the MPO pullback equal to the matrix pullback? true
Results from applying gate to MPO/ITensor/Matrix
Are all versions of computing the contraction results in the same result? true
Is the ITensor pullback equal to the matrix pullback? true
Is the MPO pullback equal to the matrix pullback? false
[-0.1423413716309932 -0.3131385545124809; -1.5979554418491506 -0.19790363455598498]
[0.7198139139346542 -1.652092462241267; 1.9081148486222423 -4.176254592261031]

with the last test failing, indicating the two pullbacks to be different.

Minimal demonstration of new behavior

After the fix one gets the new output:

Results from applying gate to MPO/ITensor/Matrix
Are all versions of computing the contraction results in the same result? true
Is the ITensor pullback equal to the matrix pullback? true
Is the MPO pullback equal to the matrix pullback? true
Results from applying gate to MPO/ITensor/Matrix
Are all versions of computing the contraction results in the same result? true
Is the ITensor pullback equal to the matrix pullback? true
Is the MPO pullback equal to the matrix pullback? true
[-0.1423413716309932 -0.3131385545124809; -1.5979554418491506 -0.19790363455598498]
[-0.1423413716309932 -0.3131385545124809; -1.5979554418491506 -0.19790363455598498]

With all test returning the desired result.

How Has This Been Tested?

I tested the behavior using the snipped in the above section. I especially run the script with all possible combinations, I.e. the field being Floats or Complex, hermitian or not hermitian and using QN numbers or not. All tests were successful.

I did not implemented the code snipped as a test in the test folder, but it could be done if wanted.

Checklist:

• [x] My code follows the style guidelines of this project. Please run using JuliaFormatter; format(".") in the base directory of the repository (~/.julia/dev/ITensors) to format your code according to our style guidelines.
• [x] I have performed a self-review of my own code.
• [x] I have commented my code, particularly in hard-to-understand areas.
• [ ] I have added tests that verify the behavior of the changes I made.
• [ ] I have made corresponding changes to the documentation.
• [x] My changes generate no new warnings.
• [ ] Any dependent changes have been merged and published in downstream modules.

[mtfishman]

[test ITensors mps]

[github-actions[bot]]

Run ITensors mps tests from comment trigger: failed ❌ https://github.com/ITensor/ITensors.jl/actions/runs/8070505033

[github-actions[bot]]

Run ITensors mps tests from comment trigger: failed ❌ https://github.com/ITensor/ITensors.jl/actions/runs/8070505033

[mtfishman]

I've already made a test for this case here: https://github.com/ITensor/ITensors.jl/blob/v0.3.57/test/ITensorLegacyMPS/ITensorChainRules/test_chainrules.jl#L172-L181, where I marked that it should throw an error but we can repurpose it now that it is working.

[ntausend]

Ah ok, was not aware of that test. Thanks!

[ntausend]

Hey, I was wondering if the PR will be included in the near future. Should I update the test in my fork and update the merge request?

[mtfishman]

Yes, please update the test to test out this new functionality.

[mtfishman]

[test ITensorMPS][test ITensorGaussianMPS]

[github-actions[bot]]

Run ITensorGaussianMPS tests from comment trigger: succeeded ✅ https://github.com/ITensor/ITensors.jl/actions/runs/8542155592

[github-actions[bot]]

Run ITensorGaussianMPS tests from comment trigger: succeeded ✅ https://github.com/ITensor/ITensors.jl/actions/runs/8542155592

[github-actions[bot]]

Run ITensorMPS tests from comment trigger: succeeded ✅ https://github.com/ITensor/ITensors.jl/actions/runs/8542155586

[github-actions[bot]]

Run ITensorMPS tests from comment trigger: succeeded ✅ https://github.com/ITensor/ITensors.jl/actions/runs/8542155586

[mtfishman]

Looks good, thanks!