ITensor / ITensorMPS.jl

MPS and MPO methods based on ITensor (ITensors.jl)
Apache License 2.0
24 stars 7 forks source link

[BUG] ChainRule for norm function #95

Open davidsmp opened 1 week ago

davidsmp commented 1 week ago

Description of bug

Using AD to differentiate the norm gives error: "MethodError: Cannot convert an object of type ChainRulesCore.Tangent{Any, @NamedTuple{data::Vector{Any}, llim::ChainRulesCore.ZeroTangent, rlim::ChainRulesCore.ZeroTangent}} to an object of type MPS"

Minimal code demonstrating the bug or unexpected behavior

Minimal runnable code

```julia using ITensors, ITensorMPS using LinearAlgebra using Zygote function randhaar(s1,s2) U = reshape(Matrix(qr(rand(ComplexF64,4,4)).Q),2,2,2,2) return ITensor(U,s1,s2,s1',s2') end N = 6 d = 2 sites = siteinds(d,N) M = randomMPS(sites,2) function F(us::Vector{ITensor}) return norm(apply(us,M)) end us = ITensor[randhaar(s1,s2) for (s1,s2) in zip(sites[1:end-1],sites[2:end])] #F(us) ≈ 1 gradient(F,us) #fails ```

Expected output or behavior

Using the function inner instead of norm works, so the problem is probably the rrule for norm

Version information

 ~/.julia/environments/v1.10/Project.toml ⁠ ⌃ [082447d4] ChainRules v1.71.0 [d360d2e6] ChainRulesCore v1.25.0 ⌃ [0d1a4710] ITensorMPS v0.3.1 ⌃ [9136182c] ITensors v0.7.3 ⌃ [e88e6eb3] Zygote v0.6.72

mtfishman commented 1 week ago

Thanks for the report. As a workaround, can you use:

function F(us::Vector{ITensor})
  x = apply(us,M)
  return sqrt(real(inner(x, x)))
end

?

mtfishman commented 1 week ago

Also, I moved the issue to the ITensorMPS.jl repository since that is where our MPS/MPO code lives now.

davidsmp commented 1 week ago

Using the inner function works, yes.

Thanks for the report. As a workaround, can you use:

function F(us::Vector{ITensor})
  x = apply(us,M)
  return sqrt(real(inner(x, x)))
end

?

mtfishman commented 1 week ago

Great, thanks for confirming. Ideally norm would work in the same way, if it was defined as sqrt(real(inner(x, x))) internally it would work, but I think we have extra logic to handle potential overflow/underflow in the case of very big or small norms which is tripping up Zygote.