[BUG] ChainRule for norm function

[davidsmp]

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

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]

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]

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

[davidsmp]

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]

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.