search

GTorlai / PastaQ.jl

Package for Simulation, Tomography and Analysis of Quantum Computers
Apache License 2.0
142 stars 23 forks source link

ChainRule error when calculating gradient through noise #300

Open danielalcalde opened 1 year ago

danielalcalde commented 1 year ago

I encountered a peculiar bug where parts of my code that utilize noise have suddenly stopped working, despite the fact that I have not updated PastaQ or any other package. This can be confirmed by referring to my Manifest file, which is tracked by git. Interestingly, rolling back to version 0.0.20 of PastaQ resolves the issue.

I'm using Julia 1.8.5 but observed the same error with julia 1.7.3

Below is a minimal example illustrating the problem, adapted from the official documentation with minor modifications to incorporate noise:

using PastaQ
using ITensors
using Zygote

N = 2   # number of qubits
J = 1.0  # Ising exchange interaction

# Hilbert space
hilbert = qubits(N)

function ising_hamiltonian(N; J)
  os = OpSum()
  for j in 1:(N - 1)
    os += -J, "Z", j, "Z", j + 1
  end
  return os
end

# define the Hamiltonian
os = ising_hamiltonian(N; J)

# build MPO "cost function"
H = MPO(os, hilbert)

noise= (1 => ("depolarizing", (p = 0.1,)), 
        2 => ("depolarizing", (p = 0.5,)))

# layer of single-qubit Ry gates
Rylayer(N, θ) = [("Ry", j, (θ=θ[j],)) for j in 1:N]

# variational ansatz
function variationalcircuit(N, depth, theta)
  circuit = Tuple[]
  for d in 1:depth
    circuit = vcat(circuit, Rylayer(N, theta[d]))
  end
  return circuit
end

depth = 1
ψ = productstate(hilbert)
ρ = outer(ψ, ψ')

cutoff = 1e-8
maxdim = 50

# cost function
function loss(theta)
  circuit = variationalcircuit(N, depth, theta)
  Uρ = runcircuit(ρ, circuit; cutoff, maxdim, noise)
  return real(inner(Uρ, H; cutoff, maxdim))
end

# initialize parameters
theta_s = [2π .* rand(N) for _ in 1:depth]

function loss_and_grad(x)
  y, (∇,) = withgradient(loss, x)
  return y, ∇
end
loss_and_grad(theta_s)

for which I get this error for version 0.0.24

MethodError: no method matching (::ChainRulesCore.ProjectTo{AbstractArray, NamedTuple{(:elements, :axes), Tuple{Vector{Any}, Tuple{Base.OneTo{Int64}}}}})(::Tuple{})
Closest candidates are:
  (::ChainRulesCore.ProjectTo{AbstractArray})(::Union{LinearAlgebra.Adjoint{T, var"#s886"}, LinearAlgebra.Transpose{T, var"#s886"}} where {T, var"#s886"<:(AbstractVector)}) at ~/.julia/packages/ChainRulesCore/C73ay/src/projection.jl:247
  (::ChainRulesCore.ProjectTo)(::ChainRulesCore.Thunk) at ~/.julia/packages/ChainRulesCore/C73ay/src/projection.jl:124
  (::ChainRulesCore.ProjectTo{AbstractArray})(::ChainRulesCore.Tangent) at /local/alcalde/.julia/packages/Zygote/dABKa/src/compiler/chainrules.jl:195
  ...

Stacktrace:
  [1] (::ChainRules.var"#1400#1406"{ChainRulesCore.ProjectTo{AbstractArray, NamedTuple{(:elements, :axes), Tuple{Vector{Any}, Tuple{Base.OneTo{Int64}}}}}, Tuple{UnitRange{Int64}}, ChainRulesCore.Tangent{Any, Tuple{Vector{ChainRulesCore.AbstractTangent}}}})()
    @ ChainRules /local/alcalde/.julia/packages/ChainRules/2ql0h/src/rulesets/Base/array.jl:307
  [2] unthunk
    @ ~/.julia/packages/ChainRulesCore/C73ay/src/tangent_types/thunks.jl:204 [inlined]
  [3] unthunk(x::ChainRulesCore.InplaceableThunk{ChainRulesCore.Thunk{ChainRules.var"#1400#1406"{ChainRulesCore.ProjectTo{AbstractArray, NamedTuple{(:elements, :axes), Tuple{Vector{Any}, Tuple{Base.OneTo{Int64}}}}}, Tuple{UnitRange{Int64}}, ChainRulesCore.Tangent{Any, Tuple{Vector{ChainRulesCore.AbstractTangent}}}}}, ChainRules.var"#1399#1405"{Tuple{UnitRange{Int64}}, ChainRulesCore.Tangent{Any, Tuple{Vector{ChainRulesCore.AbstractTangent}}}}})
    @ ChainRulesCore ~/.julia/packages/ChainRulesCore/C73ay/src/tangent_types/thunks.jl:237
  [4] wrap_chainrules_output
    @ /local/alcalde/.julia/packages/Zygote/dABKa/src/compiler/chainrules.jl:105 [inlined]
  [5] map(f::typeof(Zygote.wrap_chainrules_output), t::Tuple{ChainRulesCore.NoTangent, ChainRulesCore.InplaceableThunk{ChainRulesCore.Thunk{ChainRules.var"#1400#1406"{ChainRulesCore.ProjectTo{AbstractArray, NamedTuple{(:elements, :axes), Tuple{Vector{Any}, Tuple{Base.OneTo{Int64}}}}}, Tuple{UnitRange{Int64}}, ChainRulesCore.Tangent{Any, Tuple{Vector{ChainRulesCore.AbstractTangent}}}}}, ChainRules.var"#1399#1405"{Tuple{UnitRange{Int64}}, ChainRulesCore.Tangent{Any, Tuple{Vector{ChainRulesCore.AbstractTangent}}}}}, ChainRulesCore.InplaceableThunk{ChainRulesCore.Thunk{ChainRules.var"#1400#1406"{ChainRulesCore.ProjectTo{AbstractArray, NamedTuple{(:elements, :axes), Tuple{Vector{ChainRulesCore.ProjectTo{AbstractArray, NamedTuple{(:elements, :axes), Tuple{Vector{ChainRulesCore.ProjectTo}, Tuple{Base.OneTo{Int64}}}}}}, Tuple{Base.OneTo{Int64}}}}}, Tuple{UnitRange{Int64}}, ChainRulesCore.Tangent{Any, Tuple{Vector{ChainRulesCore.AbstractTangent}}}}}, ChainRules.var"#1399#1405"{Tuple{UnitRange{Int64}}, ChainRulesCore.Tangent{Any, Tuple{Vector{ChainRulesCore.AbstractTangent}}}}}})
    @ Base ./tuple.jl:223
  [6] wrap_chainrules_output
    @ /local/alcalde/.julia/packages/Zygote/dABKa/src/compiler/chainrules.jl:106 [inlined]
  [7] (::Zygote.ZBack{ChainRules.var"#vcat_pullback#1402"{Tuple{ChainRulesCore.ProjectTo{AbstractArray, NamedTuple{(:elements, :axes), Tuple{Vector{Any}, Tuple{Base.OneTo{Int64}}}}}, ChainRulesCore.ProjectTo{AbstractArray, NamedTuple{(:elements, :axes), Tuple{Vector{ChainRulesCore.ProjectTo{AbstractArray, NamedTuple{(:elements, :axes), Tuple{Vector{ChainRulesCore.ProjectTo}, Tuple{Base.OneTo{Int64}}}}}}, Tuple{Base.OneTo{Int64}}}}}}, Tuple{Tuple{Int64}, Tuple{Int64}}, Val{1}}})(dy::Tuple{Vector{Union{Nothing, Tuple{Nothing, Nothing, NamedTuple{(:θ,), Tuple{Float64}}}}}})
    @ Zygote /local/alcalde/.julia/packages/Zygote/dABKa/src/compiler/chainrules.jl:206
  [8] Pullback
    @ ~/.julia/packages/PastaQ/D5CCg/src/circuits/noise.jl:193 [inlined]
  [9] (::typeof(∂(#insertnoise#89)))(Δ::Tuple{Vector{Union{Nothing, Tuple{Nothing, Nothing, NamedTuple{(:θ,), Tuple{Float64}}}}}})
    @ Zygote /local/alcalde/.julia/packages/Zygote/dABKa/src/compiler/interface2.jl:0
 [10] Pullback
    @ ~/.julia/packages/PastaQ/D5CCg/src/circuits/noise.jl:141 [inlined]
 [11] (::typeof(∂(insertnoise)))(Δ::Tuple{Vector{Union{Nothing, Tuple{Nothing, Nothing, NamedTuple{(:θ,), Tuple{Float64}}}}}})
    @ Zygote /local/alcalde/.julia/packages/Zygote/dABKa/src/compiler/interface2.jl:0
 [12] Pullback
    @ ~/.julia/packages/PastaQ/D5CCg/src/circuits/noise.jl:199 [inlined]
 [13] (::typeof(∂(#insertnoise#96)))(Δ::Vector{Tuple{Nothing, Nothing, Any}})
    @ Zygote /local/alcalde/.julia/packages/Zygote/dABKa/src/compiler/interface2.jl:0
 [14] Pullback
    @ ~/.julia/packages/PastaQ/D5CCg/src/circuits/noise.jl:198 [inlined]
 [15] (::typeof(∂(insertnoise)))(Δ::Vector{Tuple{Nothing, Nothing, Any}})
    @ Zygote /local/alcalde/.julia/packages/Zygote/dABKa/src/compiler/interface2.jl:0
 [16] Pullback
    @ ~/.julia/packages/PastaQ/D5CCg/src/circuits/runcircuit.jl:24 [inlined]
 [17] (::typeof(∂(#buildcircuit#128)))(Δ::Vector{ITensor})
    @ Zygote /local/alcalde/.julia/packages/Zygote/dABKa/src/compiler/interface2.jl:0
 [18] Pullback
    @ ~/.julia/packages/PastaQ/D5CCg/src/circuits/runcircuit.jl:15 [inlined]
 [19] (::typeof(∂(buildcircuit##kw)))(Δ::Vector{ITensor})
    @ Zygote /local/alcalde/.julia/packages/Zygote/dABKa/src/compiler/interface2.jl:0
 [20] #208
    @ /local/alcalde/.julia/packages/Zygote/dABKa/src/lib/lib.jl:206 [inlined]
 [21] (::Zygote.var"#1914#back#210"{Zygote.var"#208#209"{Tuple{Tuple{Nothing, Nothing, Nothing}, Tuple{Nothing}}, typeof(∂(buildcircuit##kw))}})(Δ::Vector{ITensor})
    @ Zygote ~/.julia/packages/ZygoteRules/AIbCs/src/adjoint.jl:67
 [22] Pullback
    @ ~/.julia/packages/PastaQ/D5CCg/src/circuits/runcircuit.jl:39 [inlined]
 [23] (::typeof(∂(#buildcircuit#132)))(Δ::Vector{ITensor})
    @ Zygote /local/alcalde/.julia/packages/Zygote/dABKa/src/compiler/interface2.jl:0
 [24] (::Zygote.var"#208#209"{Tuple{Tuple{Nothing, Nothing, Nothing}, Tuple{Nothing}}, typeof(∂(#buildcircuit#132))})(Δ::Vector{ITensor})
    @ Zygote /local/alcalde/.julia/packages/Zygote/dABKa/src/lib/lib.jl:206
 [25] #1914#back
    @ ~/.julia/packages/ZygoteRules/AIbCs/src/adjoint.jl:67 [inlined]
 [26] Pullback
    @ ~/.julia/packages/PastaQ/D5CCg/src/circuits/runcircuit.jl:38 [inlined]
 [27] (::typeof(∂(buildcircuit##kw)))(Δ::Vector{ITensor})
    @ Zygote /local/alcalde/.julia/packages/Zygote/dABKa/src/compiler/interface2.jl:0
 [28] Pullback
    @ ~/.julia/packages/PastaQ/D5CCg/src/circuits/runcircuit.jl:163 [inlined]
 [29] (::typeof(∂(#runcircuit#145)))(Δ::MPO)
    @ Zygote /local/alcalde/.julia/packages/Zygote/dABKa/src/compiler/interface2.jl:0
 [30] Pullback
    @ ~/.julia/packages/PastaQ/D5CCg/src/circuits/runcircuit.jl:153 [inlined]
 [31] (::typeof(∂(runcircuit##kw)))(Δ::MPO)
    @ Zygote /local/alcalde/.julia/packages/Zygote/dABKa/src/compiler/interface2.jl:0
 [32] Pullback
    @ ./In[7]:23 [inlined]
 [33] (::typeof(∂(loss)))(Δ::Float64)
    @ Zygote /local/alcalde/.julia/packages/Zygote/dABKa/src/compiler/interface2.jl:0
 [34] (::Zygote.var"#60#61"{typeof(∂(loss))})(Δ::Float64)
    @ Zygote /local/alcalde/.julia/packages/Zygote/dABKa/src/compiler/interface.jl:45
 [35] withgradient(f::Function, args::Vector{Vector{Float64}})
    @ Zygote /local/alcalde/.julia/packages/Zygote/dABKa/src/compiler/interface.jl:124
 [36] loss_and_grad(x::Vector{Vector{Float64}})
    @ Main ./In[7]:31
 [37] top-level scope
    @ In[7]:34
 [38] eval
    @ ./boot.jl:368 [inlined]
 [39] include_string(mapexpr::typeof(REPL.softscope), mod::Module, code::String, filename::String)
    @ Base ./loading.jl:1428

and for versions 0.0.21-0.0.23 I get the error:

MethodError: reducing over an empty collection is not allowed; consider supplying `init` to the reducer

Stacktrace:
  [1] macro expansion
    @ ~/.julia/packages/Zygote/PD12J/src/compiler/interface2.jl:0 [inlined]
  [2] _pullback(::Zygote.Context{false}, ::typeof(Base.reduce_empty), ::typeof(max), ::Type{Union{}})
    @ Zygote ~/.julia/packages/Zygote/PD12J/src/compiler/interface2.jl:9
  [3] _pullback
    @ ./reduce.jl:355 [inlined]
  [4] _pullback(::Zygote.Context{false}, ::typeof(Base.reduce_empty), ::Base.BottomRF{typeof(max)}, ::Type{Union{}})
    @ Zygote ~/.julia/packages/Zygote/PD12J/src/compiler/interface2.jl:0
  [5] _pullback
    @ ./reduce.jl:379 [inlined]
  [6] _pullback
    @ ./reduce.jl:378 [inlined]
  [7] _pullback
    @ ./reduce.jl:49 [inlined]
  [8] _pullback(::Zygote.Context{false}, ::typeof(Base.foldl_impl), ::Base.BottomRF{typeof(max)}, ::Base._InitialValue, ::Tuple{})
    @ Zygote ~/.julia/packages/Zygote/PD12J/src/compiler/interface2.jl:0
  [9] _pullback
    @ ./reduce.jl:44 [inlined]
 [10] _pullback(::Zygote.Context{false}, ::typeof(Base.mapfoldl_impl), ::typeof(identity), ::typeof(max), ::Base._InitialValue, ::Tuple{})
    @ Zygote ~/.julia/packages/Zygote/PD12J/src/compiler/interface2.jl:0
 [11] _pullback
    @ ./reduce.jl:170 [inlined]
--- the last 2 lines are repeated 1 more time ---
 [14] _pullback(::Zygote.Context{false}, ::typeof(mapfoldl), ::typeof(identity), ::typeof(max), ::Tuple{})
    @ Zygote ~/.julia/packages/Zygote/PD12J/src/compiler/interface2.jl:0
 [15] _pullback
    @ ./reduce.jl:302 [inlined]
--- the last 2 lines are repeated 1 more time ---
 [18] _pullback(::Zygote.Context{false}, ::typeof(mapreduce), ::typeof(identity), ::typeof(max), ::Tuple{})
    @ Zygote ~/.julia/packages/Zygote/PD12J/src/compiler/interface2.jl:0
 [19] _pullback
    @ ./reduce.jl:757 [inlined]
--- the last 2 lines are repeated 1 more time ---
 [22] _pullback(ctx::Zygote.Context{false}, f::typeof(maximum), args::Tuple{})
    @ Zygote ~/.julia/packages/Zygote/PD12J/src/compiler/interface2.jl:0
 [23] _pullback
    @ /local/alcalde/.julia/packages/PastaQ/9DfPk/src/circuits/circuits.jl:4 [inlined]
 [24] _pullback(ctx::Zygote.Context{false}, f::typeof(nqubits), args::Tuple{String, Int64, NamedTuple{(:θ,), Tuple{Float64}}})
    @ Zygote ~/.julia/packages/Zygote/PD12J/src/compiler/interface2.jl:0
 [25] _pullback
    @ ./none:0 [inlined]
 [26] _pullback(ctx::Zygote.Context{false}, f::PastaQ.var"#106#107", args::Tuple{String, Int64, NamedTuple{(:θ,), Tuple{Float64}}})
    @ Zygote ~/.julia/packages/Zygote/PD12J/src/compiler/interface2.jl:0
 [27] _pullback
    @ ./reduce.jl:95 [inlined]
 [28] _pullback(::Zygote.Context{false}, ::Base.MappingRF{PastaQ.var"#106#107", Base.BottomRF{typeof(max)}}, ::Base._InitialValue, ::Tuple{String, Int64, NamedTuple{(:θ,), Tuple{Float64}}})
    @ Zygote ~/.julia/packages/Zygote/PD12J/src/compiler/interface2.jl:0
 [29] _pullback
    @ ./reduce.jl:58 [inlined]
 [30] _pullback(::Zygote.Context{false}, ::typeof(Base._foldl_impl), ::Base.MappingRF{PastaQ.var"#106#107", Base.BottomRF{typeof(max)}}, ::Base._InitialValue, ::Vector{Tuple})
    @ Zygote ~/.julia/packages/Zygote/PD12J/src/compiler/interface2.jl:0
 [31] _pullback
    @ ./reduce.jl:48 [inlined]
 [32] _pullback(::Zygote.Context{false}, ::typeof(Base.foldl_impl), ::Base.MappingRF{PastaQ.var"#106#107", Base.BottomRF{typeof(max)}}, ::Base._InitialValue, ::Vector{Tuple})
    @ Zygote ~/.julia/packages/Zygote/PD12J/src/compiler/interface2.jl:0
 [33] _pullback
    @ ./reduce.jl:44 [inlined]
 [34] _pullback(::Zygote.Context{false}, ::typeof(Base.mapfoldl_impl), ::typeof(identity), ::typeof(max), ::Base._InitialValue, ::Base.Generator{Vector{Tuple}, PastaQ.var"#106#107"})
    @ Zygote ~/.julia/packages/Zygote/PD12J/src/compiler/interface2.jl:0
 [35] _pullback
    @ ./reduce.jl:170 [inlined]
--- the last 2 lines are repeated 1 more time ---
 [38] _pullback(::Zygote.Context{false}, ::typeof(mapfoldl), ::typeof(identity), ::typeof(max), ::Base.Generator{Vector{Tuple}, PastaQ.var"#106#107"})
    @ Zygote ~/.julia/packages/Zygote/PD12J/src/compiler/interface2.jl:0
 [39] _pullback (repeats 2 times)
    @ ./reduce.jl:302 [inlined]
 [40] _pullback (repeats 2 times)
    @ ./reduce.jl:757 [inlined]
 [41] _pullback
    @ /local/alcalde/.julia/packages/PastaQ/9DfPk/src/circuits/circuits.jl:6 [inlined]
 [42] _pullback
    @ ./none:0 [inlined]
 [43] _pullback
    @ ./reduce.jl:95 [inlined]
 [44] _pullback
    @ ./reduce.jl:58 [inlined]
 [45] _pullback(::Zygote.Context{false}, ::typeof(Base._foldl_impl), ::Base.MappingRF{PastaQ.var"#106#107", Base.BottomRF{typeof(max)}}, ::Base._InitialValue, ::Vector{Vector{Tuple}})
    @ Zygote ~/.julia/packages/Zygote/PD12J/src/compiler/interface2.jl:0
 [46] _pullback
    @ ./reduce.jl:48 [inlined]
 [47] _pullback(::Zygote.Context{false}, ::typeof(Base.foldl_impl), ::Base.MappingRF{PastaQ.var"#106#107", Base.BottomRF{typeof(max)}}, ::Base._InitialValue, ::Vector{Vector{Tuple}})
    @ Zygote ~/.julia/packages/Zygote/PD12J/src/compiler/interface2.jl:0
 [48] _pullback
    @ ./reduce.jl:44 [inlined]
 [49] _pullback(::Zygote.Context{false}, ::typeof(Base.mapfoldl_impl), ::typeof(identity), ::typeof(max), ::Base._InitialValue, ::Base.Generator{Vector{Vector{Tuple}}, PastaQ.var"#106#107"})
    @ Zygote ~/.julia/packages/Zygote/PD12J/src/compiler/interface2.jl:0
 [50] _pullback
    @ ./reduce.jl:170 [inlined]
--- the last 2 lines are repeated 1 more time ---
 [53] _pullback(::Zygote.Context{false}, ::typeof(mapfoldl), ::typeof(identity), ::typeof(max), ::Base.Generator{Vector{Vector{Tuple}}, PastaQ.var"#106#107"})
    @ Zygote ~/.julia/packages/Zygote/PD12J/src/compiler/interface2.jl:0
 [54] _pullback (repeats 2 times)
    @ ./reduce.jl:302 [inlined]
 [55] _pullback (repeats 2 times)
    @ ./reduce.jl:757 [inlined]
 [56] _pullback
    @ /local/alcalde/.julia/packages/PastaQ/9DfPk/src/circuits/circuits.jl:6 [inlined]
 [57] _pullback
    @ /local/alcalde/.julia/packages/PastaQ/9DfPk/src/circuits/noise.jl:142 [inlined]
 [58] _pullback(::Zygote.Context{false}, ::PastaQ.var"##insertnoise#89", ::Nothing, ::typeof(insertnoise), ::Vector{Vector{Tuple}}, ::Tuple{Pair{Int64, Tuple{String, NamedTuple{(:p,), Tuple{Float64}}}}, Pair{Int64, Tuple{String, NamedTuple{(:p,), Tuple{Float64}}}}})
    @ Zygote ~/.julia/packages/Zygote/PD12J/src/compiler/interface2.jl:0
 [59] _pullback
    @ /local/alcalde/.julia/packages/PastaQ/9DfPk/src/circuits/noise.jl:140 [inlined]
 [60] _pullback
    @ /local/alcalde/.julia/packages/PastaQ/9DfPk/src/circuits/noise.jl:198 [inlined]
 [61] _pullback(::Zygote.Context{false}, ::PastaQ.var"##insertnoise#96", ::Base.Pairs{Symbol, Union{}, Tuple{}, NamedTuple{(), Tuple{}}}, ::typeof(insertnoise), ::Vector{Tuple}, ::Tuple{Pair{Int64, Tuple{String, NamedTuple{(:p,), Tuple{Float64}}}}, Pair{Int64, Tuple{String, NamedTuple{(:p,), Tuple{Float64}}}}})
    @ Zygote ~/.julia/packages/Zygote/PD12J/src/compiler/interface2.jl:0
 [62] _pullback
    @ /local/alcalde/.julia/packages/PastaQ/9DfPk/src/circuits/noise.jl:197 [inlined]
 [63] _pullback
    @ /local/alcalde/.julia/packages/PastaQ/9DfPk/src/circuits/runcircuit.jl:24 [inlined]
 [64] _pullback(::Zygote.Context{false}, ::PastaQ.var"##buildcircuit#128", ::Tuple{Pair{Int64, Tuple{String, NamedTuple{(:p,), Tuple{Float64}}}}, Pair{Int64, Tuple{String, NamedTuple{(:p,), Tuple{Float64}}}}}, ::Nothing, ::typeof(identity), ::typeof(buildcircuit), ::Vector{Index{Int64}}, ::Vector{Tuple})
    @ Zygote ~/.julia/packages/Zygote/PD12J/src/compiler/interface2.jl:0
 [65] _pullback
    @ /local/alcalde/.julia/packages/PastaQ/9DfPk/src/circuits/runcircuit.jl:15 [inlined]
 [66] _pullback(::Zygote.Context{false}, ::PastaQ.var"#buildcircuit##kw", ::NamedTuple{(:noise, :eltype, :device), Tuple{Tuple{Pair{Int64, Tuple{String, NamedTuple{(:p,), Tuple{Float64}}}}, Pair{Int64, Tuple{String, NamedTuple{(:p,), Tuple{Float64}}}}}, Nothing, typeof(identity)}}, ::typeof(buildcircuit), ::Vector{Index{Int64}}, ::Vector{Tuple})
    @ Zygote ~/.julia/packages/Zygote/PD12J/src/compiler/interface2.jl:0
 [67] _apply(::Function, ::Vararg{Any})
    @ Core ./boot.jl:816
 [68] adjoint
    @ ~/.julia/packages/Zygote/PD12J/src/lib/lib.jl:203 [inlined]
 [69] _pullback
    @ ~/.julia/packages/ZygoteRules/AIbCs/src/adjoint.jl:65 [inlined]
 [70] _pullback
    @ /local/alcalde/.julia/packages/PastaQ/9DfPk/src/circuits/runcircuit.jl:39 [inlined]
 [71] _pullback(::Zygote.Context{false}, ::PastaQ.var"##buildcircuit#132", ::Base.Pairs{Symbol, Any, Tuple{Symbol, Symbol, Symbol}, NamedTuple{(:noise, :eltype, :device), Tuple{Tuple{Pair{Int64, Tuple{String, NamedTuple{(:p,), Tuple{Float64}}}}, Pair{Int64, Tuple{String, NamedTuple{(:p,), Tuple{Float64}}}}}, Nothing, typeof(identity)}}}, ::typeof(buildcircuit), ::MPO, ::Vector{Tuple})
    @ Zygote ~/.julia/packages/Zygote/PD12J/src/compiler/interface2.jl:0
 [72] _apply(::Function, ::Vararg{Any})
    @ Core ./boot.jl:816
 [73] adjoint
    @ ~/.julia/packages/Zygote/PD12J/src/lib/lib.jl:203 [inlined]
 [74] _pullback
    @ ~/.julia/packages/ZygoteRules/AIbCs/src/adjoint.jl:65 [inlined]
 [75] _pullback
    @ /local/alcalde/.julia/packages/PastaQ/9DfPk/src/circuits/runcircuit.jl:38 [inlined]
 [76] _pullback
    @ /local/alcalde/.julia/packages/PastaQ/9DfPk/src/circuits/runcircuit.jl:163 [inlined]
 [77] _pullback(::Zygote.Context{false}, ::PastaQ.var"##runcircuit#145", ::Bool, ::Tuple{Pair{Int64, Tuple{String, NamedTuple{(:p,), Tuple{Float64}}}}, Pair{Int64, Tuple{String, NamedTuple{(:p,), Tuple{Float64}}}}}, ::Nothing, ::typeof(identity), ::Base.Pairs{Symbol, Real, Tuple{Symbol, Symbol}, NamedTuple{(:cutoff, :maxdim), Tuple{Float64, Int64}}}, ::typeof(runcircuit), ::MPO, ::Vector{Tuple})
    @ Zygote ~/.julia/packages/Zygote/PD12J/src/compiler/interface2.jl:0
 [78] _pullback
    @ /local/alcalde/.julia/packages/PastaQ/9DfPk/src/circuits/runcircuit.jl:153 [inlined]
 [79] _pullback(::Zygote.Context{false}, ::PastaQ.var"#runcircuit##kw", ::NamedTuple{(:cutoff, :maxdim, :noise), Tuple{Float64, Int64, Tuple{Pair{Int64, Tuple{String, NamedTuple{(:p,), Tuple{Float64}}}}, Pair{Int64, Tuple{String, NamedTuple{(:p,), Tuple{Float64}}}}}}}, ::typeof(runcircuit), ::MPO, ::Vector{Tuple})
    @ Zygote ~/.julia/packages/Zygote/PD12J/src/compiler/interface2.jl:0
 [80] _pullback
    @ ./In[9]:23 [inlined]
 [81] _pullback(ctx::Zygote.Context{false}, f::typeof(loss), args::Vector{Vector{Float64}})
    @ Zygote ~/.julia/packages/Zygote/PD12J/src/compiler/interface2.jl:0
 [82] pullback(f::Function, cx::Zygote.Context{false}, args::Vector{Vector{Float64}})
    @ Zygote ~/.julia/packages/Zygote/PD12J/src/compiler/interface.jl:44
 [83] pullback
    @ ~/.julia/packages/Zygote/PD12J/src/compiler/interface.jl:42 [inlined]
 [84] withgradient(f::Function, args::Vector{Vector{Float64}})
    @ Zygote ~/.julia/packages/Zygote/PD12J/src/compiler/interface.jl:123
 [85] loss_and_grad(x::Vector{Vector{Float64}})
    @ Main ./In[9]:31
 [86] top-level scope
    @ In[9]:34
 [87] eval
    @ ./boot.jl:368 [inlined]

Installed packages:

  [7d9f7c33] Accessors v0.1.26
⌅ [4fba245c] ArrayInterface v5.0.8
  [198e06fe] BangBang v0.3.37
  [6e4b80f9] BenchmarkTools v1.3.2
  [c3b6d118] BitIntegers v0.2.7
  [a74b3585] Blosc v0.7.3
⌅ [052768ef] CUDA v3.13.1
  [3da002f7] ColorTypes v0.11.4
  [864edb3b] DataStructures v0.18.13
  [ffbed154] DocStringExtensions v0.9.3
  [5789e2e9] FileIO v1.16.0
⌃ [587475ba] Flux v0.12.10
  [41a02a25] Folds v0.2.8
  [46192b85] GPUArraysCore v0.1.4
  [61eb1bfa] GPUCompiler v0.17.2
⌅ [f67ccb44] HDF5 v0.15.7
  [7073ff75] IJulia v1.24.0
  [9136182c] ITensors v0.3.27
  [033835bb] JLD2 v0.4.30
  [0b1a1467] KrylovKit v0.6.0
  [929cbde3] LLVM v4.16.0
  [e89f7d12] Media v0.5.0
  [85b6ec6f] MethodAnalysis v0.4.11
  [23ae76d9] NDTensors v0.1.47
  [872c559c] NNlib v0.8.18
  [77e91f04] OptimKit v0.3.1
  [30b07047] PastaQ v0.0.24
  [91a5bcdd] Plots v1.38.5
  [e6cf234a] RandomNumbers v1.5.3
  [ae029012] Requires v1.3.0
  [295af30f] Revise v3.5.1
  [aa65fe97] SnoopCompile v2.10.2
  [90137ffa] StaticArrays v1.5.15
  [1e83bf80] StaticArraysCore v1.4.0
  [5e0ebb24] Strided v1.2.3
  [a759f4b9] TimerOutputs v0.5.22
  [28d57a85] Transducers v0.4.75
  [9d95972d] TupleTools v1.3.0
  [e88e6eb3] Zygote v0.6.55
  [56f22d72] Artifacts
  [37e2e46d] LinearAlgebra
  [44cfe95a] Pkg v1.8.0

Hope someone can make sense of this.

mtfishman commented 1 year ago

That sounds really strange. The compat entries of PastaQ.jl have barely changed between those versions: https://github.com/GTorlai/PastaQ.jl/blob/v0.0.20/Project.toml https://github.com/GTorlai/PastaQ.jl/blob/v0.0.24/Project.toml

so I don't see why you would be seeing differences between those versions, since if something changed in ChainRules or Zygote, PastaQ v0.0.20 and v0.0.24 should be using the same versions of ChainRules or Zygote. From the error messages it looks like a Zygote or ChainRules issue. I have seen regressions in some Zygote operations so I could believe that an operations that we are using in PastaQ used to be differentiable but may not in a or earlier version of Zygote. Are you certain that in the comparison between the cases that work and don't work that you are using the same versions of Zygote and ChainRules? I would try using a clean project environment and doing a minimal installation of the package you need to run your script (PastaQ, ITensors, and Zygote), make sure everything is updated, and then compare between PastaQ v0.0.20 and PastaQ v0.0.24 (and check the versions of ChainRules and Zygote that get installed), as a sanity check.

Another possibility is that we are using some operations in PastaQ v0.0.24 that are no longer differentiable in recent versions of ChainRules/Zygote, but in PastaQ v0.0.20 we aren't using those operations. So you may be detecting a regression in ChainRules/Zygote, and the solution from our side or your side would be to pin ChainRules or Zygote to an older version where PastaQ v0.0.24 works (or detect which operations had a regression and rewrite those parts of the code to be differentiable with the latest versions of ChainRules/Zygote).

danielalcalde commented 1 year ago

I have installed PastaQ in two different projects, one with V0.0.20 and another with v0.0.24. Both projects have identical Manifest files, except for the PastaQ entry and the same behavior I wrote about above is happening. As a sanity check I did the same on another machine with the same effect. This makes me think that the issue is not related to Zygote.

I also tried implementing the ChainRulesCore.ProjectTo function, but it only resulted in more errors.

Implementation:

function (b::ChainRulesCore.ProjectTo{AbstractArray, NamedTuple{(:elements, :axes), Tuple{Vector{Any}, Tuple{Base.OneTo{Int64}}}}})(a::Tuple{})
    return Array{Number}(collect(a))
end

after which it wants me to implement

function (b::ChainRulesCore.ProjectTo{AbstractArray, NamedTuple{(:elements, :axes), Tuple{Vector{ChainRulesCore.ProjectTo{AbstractArray, NamedTuple{(:elements, :axes), Tuple{Vector{ChainRulesCore.ProjectTo}, Tuple{Base.OneTo{Int64}}}}}}, Tuple{Base.OneTo{Int64}}}}})(a::Tuple{Vector{ChainRulesCore.AbstractTangent}})

which honestly I do not know what it is expecting here.

mtfishman commented 1 year ago

Alright, so as I expected this is caused by a regression in Zygote. I tried running your code but pinning Zygote to an older version (Zygote v0.6.43, while the latest version is v0.6.60):

julia> ] add Zygote@0.6.43

and then running your code works.

I see the error message is originating from this line: https://github.com/GTorlai/PastaQ.jl/blob/v0.0.24/src/circuits/noise.jl#L193

So somewhere between Zygote v0.6.43 and v0.6.60, there was a regression in differentiating either that line, or some combination of code in the insertnoise function that it's in. The insertnoise is quite complicated so I'm not surprised Zygote might be having trouble differentiating it.

Anyway, that's a minimal fix that you can try from your end, let me know if that works. It would be good to know at what version Zygote starting failing to differentiate the insertnoise function, and then try to raise an issue over at Zygote about it.

mtfishman commented 1 year ago

Alright, some more information. It looks like Zygote versions v0.6.44 and below work, and v0.6.45 and above fail at this same line of code in insertnoise.

In the release notes of Zygote v0.6.45: https://github.com/FluxML/Zygote.jl/releases/tag/v0.6.45 I see this PR was merged and included in Zygote v0.6.45: https://github.com/FluxML/Zygote.jl/pull/1277 that switched the implementation of the derivative rules for hvcat, hcat, and vcat over to ChainRules. I would guess that is causing the issue in differentiating insertnoise, since it makes use of vcat.

danielalcalde commented 1 year ago

Yes, you are right, this fixes it. Thank you for the help! I wonder in which way vcat was broken since I used it in my code without issue. If I have time this week I will investigate insertnoise in more detail to find the unsupported operation.

mtfishman commented 1 year ago

Thanks! It would be nice to try to make as minimal of a code example as possible to reproduce the bug so then we can report it to Zygote. For example, a starting place could be writing a function that just depends on insertnoise but doesn't do the circuit evolution, and then try to differentiate that with Zygote. It could be a function where you make a noisy circuit with a certain gate angle and then just extract that gate angle back out of the function. Then we can try to remove any PastaQ-specific functions or types from insertnoise, which would make it easier to report to Zygote. That shouldn't be too hard do since it is basically just manipulating Vectors of Tuples. I've seen similar kinds of regressions from Zygote before, and it can be quite subtle, for example some specific combination of operations, operations that involve some control flow, etc.