[ITensors] [BUG] AD failing for delta tensors

Description of bug

Get an error when trying to differentiate a contraction of tensors involving a delta--tensor to match indices of tensors.

Minimal code demonstrating the bug or unexpected behavior

Minimal runnable code

using ITensors
using ChainRulesCore
using Zygote

function f(x, y)
    d = delta(inds(x)[1], inds(y)[1])
    return (x * d * y)[]
end

idx_1 = Index(2)
idx_2 = Index(3)
tens_1 = randomITensor(idx_1)
tens_2 = randomITensor(idx_2)
println(f(tens_1, tens_2))
l = x -> f(x, tens_2)
@assert l(tens_1) == f(tens_1, tens_2)
l'(tens_1)

Expected output or behavior

Expect a valid derivative

Actual output or behavior

Weird bounds error.

Output of minimal runnable code

BoundsError: attempt to access 2-element Vector{Float64} at index [3]

Stacktrace:
  [1] getindex
    @ ./array.jl:861 [inlined]
  [2] getindex
    @ ~/.julia/packages/NDTensors/y9JPS/src/dense/dense.jl:245 [inlined]
  [3] getindex
    @ ~/.julia/packages/NDTensors/y9JPS/src/dense/dense.jl:250 [inlined]
  [4] getdiagindex
    @ ~/.julia/packages/NDTensors/y9JPS/src/tensor.jl:346 [inlined]
  [5] contract!(C::NDTensors.DenseTensor{Float64, 1, Tuple{Index{Int64}}, NDTensors.Dense{Float64, Vector{Float64}}}, Clabels::Tuple{Int64}, A::NDTensors.DiagTensor{Float64, 2, Tuple{Index{Int64}, Index{Int64}}, NDTensors.Diag{Float64, Float64}}, Alabels::Tuple{Int64, Int64}, B::NDTensors.DenseTensor{Float64, 1, Tuple{Index{Int64}}, NDTensors.Dense{Float64, Vector{Float64}}}, Blabels::Tuple{Int64}, α::Float64, β::Float64; convert_to_dense::Bool)
    @ NDTensors ~/.julia/packages/NDTensors/y9JPS/src/diag/diag.jl:470
  [6] contract!
    @ ~/.julia/packages/NDTensors/y9JPS/src/diag/diag.jl:452 [inlined]
  [7] contract! (repeats 2 times)
    @ ~/.julia/packages/NDTensors/y9JPS/src/diag/diag.jl:562 [inlined]
  [8] _contract!!
    @ ~/.julia/packages/NDTensors/y9JPS/src/generic_tensor_operations.jl:80 [inlined]
  [9] _contract!!
    @ ~/.julia/packages/NDTensors/y9JPS/src/generic_tensor_operations.jl:77 [inlined]
 [10] contract!!
    @ ~/.julia/packages/NDTensors/y9JPS/src/generic_tensor_operations.jl:111 [inlined]
 [11] contract!!
    @ ~/.julia/packages/NDTensors/y9JPS/src/generic_tensor_operations.jl:89 [inlined]
 [12] contract(T1::NDTensors.DenseTensor{Float64, 1, Tuple{Index{Int64}}, NDTensors.Dense{Float64, Vector{Float64}}}, labelsT1::Tuple{Int64}, T2::NDTensors.DiagTensor{Float64, 2, Tuple{Index{Int64}, Index{Int64}}, NDTensors.Diag{Float64, Float64}}, labelsT2::Tuple{Int64, Int64}, labelsR::Tuple{Int64})
    @ NDTensors ~/.julia/packages/NDTensors/y9JPS/src/generic_tensor_operations.jl:69
 [13] contract(::Type{NDTensors.CanContract{NDTensors.DenseTensor{Float64, 1, Tuple{Index{Int64}}, NDTensors.Dense{Float64, Vector{Float64}}}, NDTensors.DiagTensor{Float64, 2, Tuple{Index{Int64}, Index{Int64}}, NDTensors.Diag{Float64, Float64}}}}, t1::NDTensors.DenseTensor{Float64, 1, Tuple{Index{Int64}}, NDTensors.Dense{Float64, Vector{Float64}}}, labels_t1::Tuple{Int64}, t2::NDTensors.DiagTensor{Float64, 2, Tuple{Index{Int64}, Index{Int64}}, NDTensors.Diag{Float64, Float64}}, labels_t2::Tuple{Int64, Int64})
    @ NDTensors ~/.julia/packages/NDTensors/y9JPS/src/generic_tensor_operations.jl:51
 [14] contract
    @ ~/.julia/packages/SimpleTraits/l1ZsK/src/SimpleTraits.jl:331 [inlined]
 [15] _contract(A::NDTensors.DenseTensor{Float64, 1, Tuple{Index{Int64}}, NDTensors.Dense{Float64, Vector{Float64}}}, B::NDTensors.DiagTensor{Float64, 2, Tuple{Index{Int64}, Index{Int64}}, NDTensors.Diag{Float64, Float64}})
    @ ITensors ~/.julia/packages/ITensors/PisGD/src/itensor.jl:1870
 [16] _contract(A::ITensor, B::ITensor)
    @ ITensors ~/.julia/packages/ITensors/PisGD/src/itensor.jl:1876
 [17] contract(A::ITensor, B::ITensor)
    @ ITensors ~/.julia/packages/ITensors/PisGD/src/itensor.jl:1980
 [18] *
    @ ~/.julia/packages/ITensors/PisGD/src/itensor.jl:1967 [inlined]
 [19] (::ITensors.ITensorChainRules.var"#contract_pullback#62"{ITensor, ITensor, ProjectTo{ITensor, NamedTuple{(:element,), Tuple{ProjectTo{Float64, NamedTuple{(), Tuple{}}}}}}, ProjectTo{ITensor, NamedTuple{(:element,), Tuple{ProjectTo{Float64, NamedTuple{(), Tuple{}}}}}}})(ȳ::ITensor)
    @ ITensors.ITensorChainRules ~/.julia/packages/ITensors/PisGD/src/ITensorChainRules/itensor.jl:98
 [20] (::Zygote.ZBack{ITensors.ITensorChainRules.var"#contract_pullback#62"{ITensor, ITensor, ProjectTo{ITensor, NamedTuple{(:element,), Tuple{ProjectTo{Float64, NamedTuple{(), Tuple{}}}}}}, ProjectTo{ITensor, NamedTuple{(:element,), Tuple{ProjectTo{Float64, NamedTuple{(), Tuple{}}}}}}}})(dy::ITensor)
    @ Zygote ~/.julia/packages/Zygote/DRjAT/src/compiler/chainrules.jl:206
 [21] Pullback
    @ ~/.julia/packages/ITensors/PisGD/src/itensor.jl:2066 [inlined]
 [22] (::typeof(∂(#233)))(Δ::ITensor)
    @ Zygote ~/.julia/packages/Zygote/DRjAT/src/compiler/interface2.jl:0
 [23] (::Zygote.var"#ad_pullback#50"{Tuple{ITensors.var"#233#235"{Base.Pairs{Symbol, Union{}, Tuple{}, NamedTuple{(), Tuple{}}}}, ITensor, ITensor}, typeof(∂(#233))})(Δ::ITensor)
    @ Zygote ~/.julia/packages/Zygote/DRjAT/src/compiler/chainrules.jl:258
 [24] #1679
    @ ~/.julia/packages/ChainRules/fgVxV/src/rulesets/Base/mapreduce.jl:453 [inlined]
 [25] BottomRF
    @ ./reduce.jl:81 [inlined]
 [26] (::Base.var"#789#790"{Base.BottomRF{ChainRules.var"#1679#1682"}})(::Tuple{Tuple{Tuple{NoTangent, ITensor, ITensor}}, Tuple{NoTangent, ITensor, ITensor}}, x::Tuple{ITensor, Zygote.var"#ad_pullback#50"{Tuple{ITensors.var"#233#235"{Base.Pairs{Symbol, Union{}, Tuple{}, NamedTuple{(), Tuple{}}}}, ITensor, ITensor}, typeof(∂(#233))}})
    @ Base ./accumulate.jl:310
 [27] afoldl(::Base.var"#789#790"{Base.BottomRF{ChainRules.var"#1679#1682"}}, ::Tuple{Tuple{}, Tuple{Int64, ITensor, Int64}}, ::Tuple{ITensor, Zygote.var"#ad_pullback#50"{Tuple{ITensors.var"#233#235"{Base.Pairs{Symbol, Union{}, Tuple{}, NamedTuple{(), Tuple{}}}}, ITensor, ITensor}, typeof(∂(#233))}}, ::Tuple{ITensor, Zygote.var"#ad_pullback#50"{Tuple{ITensors.var"#233#235"{Base.Pairs{Symbol, Union{}, Tuple{}, NamedTuple{(), Tuple{}}}}, ITensor, ITensor}, typeof(∂(#233))}})
    @ Base ./operators.jl:613
 [28] #accumulate#788
    @ ./accumulate.jl:309 [inlined]
 [29] (::ChainRules.var"#unfoldl#1681"{ChainRules._InitialValue, ProjectTo{Tangent{Tuple{ITensor, ITensor, ITensor}}, NamedTuple{(:elements,), Tuple{Tuple{ProjectTo{ITensor, NamedTuple{(:element,), Tuple{ProjectTo{Float64, NamedTuple{(), Tuple{}}}}}}, ProjectTo{ITensor, NamedTuple{(:element,), Tuple{ProjectTo{Float64, NamedTuple{(), Tuple{}}}}}}, ProjectTo{ITensor, NamedTuple{(:element,), Tuple{ProjectTo{Float64, NamedTuple{(), Tuple{}}}}}}}}}}, Tuple{Base.OneTo{Int64}}, Tuple{Tuple{ITensor, Zygote.var"#ad_pullback#50"{Tuple{ITensors.var"#233#235"{Base.Pairs{Symbol, Union{}, Tuple{}, NamedTuple{(), Tuple{}}}}, ITensor, ITensor}, typeof(∂(#233))}}, Tuple{ITensor, Zygote.var"#ad_pullback#50"{Tuple{ITensors.var"#233#235"{Base.Pairs{Symbol, Union{}, Tuple{}, NamedTuple{(), Tuple{}}}}, ITensor, ITensor}, typeof(∂(#233))}}}})(dy::ITensor)
    @ ChainRules ~/.julia/packages/ChainRules/fgVxV/src/rulesets/Base/mapreduce.jl:452
 [30] (::Zygote.ZBack{ChainRules.var"#unfoldl#1681"{ChainRules._InitialValue, ProjectTo{Tangent{Tuple{ITensor, ITensor, ITensor}}, NamedTuple{(:elements,), Tuple{Tuple{ProjectTo{ITensor, NamedTuple{(:element,), Tuple{ProjectTo{Float64, NamedTuple{(), Tuple{}}}}}}, ProjectTo{ITensor, NamedTuple{(:element,), Tuple{ProjectTo{Float64, NamedTuple{(), Tuple{}}}}}}, ProjectTo{ITensor, NamedTuple{(:element,), Tuple{ProjectTo{Float64, NamedTuple{(), Tuple{}}}}}}}}}}, Tuple{Base.OneTo{Int64}}, Tuple{Tuple{ITensor, Zygote.var"#ad_pullback#50"{Tuple{ITensors.var"#233#235"{Base.Pairs{Symbol, Union{}, Tuple{}, NamedTuple{(), Tuple{}}}}, ITensor, ITensor}, typeof(∂(#233))}}, Tuple{ITensor, Zygote.var"#ad_pullback#50"{Tuple{ITensors.var"#233#235"{Base.Pairs{Symbol, Union{}, Tuple{}, NamedTuple{(), Tuple{}}}}, ITensor, ITensor}, typeof(∂(#233))}}}}})(dy::ITensor)
    @ Zygote ~/.julia/packages/Zygote/DRjAT/src/compiler/chainrules.jl:206
 [31] Pullback
    @ ~/.julia/packages/ITensors/PisGD/src/itensor.jl:2066 [inlined]
 [32] (::typeof(∂(#contract#232)))(Δ::ITensor)
    @ Zygote ~/.julia/packages/Zygote/DRjAT/src/compiler/interface2.jl:0
 [33] Pullback
    @ ~/.julia/packages/ITensors/PisGD/src/itensor.jl:2065 [inlined]
 [34] (::typeof(∂(contract)))(Δ::ITensor)
    @ Zygote ~/.julia/packages/Zygote/DRjAT/src/compiler/interface2.jl:0
 [35] Pullback
    @ ~/.julia/packages/ITensors/PisGD/src/itensor.jl:2076 [inlined]
 [36] (::typeof(∂(#contract#237)))(Δ::ITensor)
    @ Zygote ~/.julia/packages/Zygote/DRjAT/src/compiler/interface2.jl:0
 [37] (::Zygote.var"#208#209"{Tuple{Tuple{Nothing, Nothing}, Tuple{Nothing, Nothing, Nothing}}, typeof(∂(#contract#237))})(Δ::ITensor)
    @ Zygote ~/.julia/packages/Zygote/DRjAT/src/lib/lib.jl:206
 [38] #1912#back
    @ ~/.julia/packages/ZygoteRules/AIbCs/src/adjoint.jl:67 [inlined]
 [39] Pullback
    @ ~/.julia/packages/ITensors/PisGD/src/itensor.jl:2076 [inlined]
 [40] (::typeof(∂(contract)))(Δ::ITensor)
    @ Zygote ~/.julia/packages/Zygote/DRjAT/src/compiler/interface2.jl:0
 [41] #208
    @ ~/.julia/packages/Zygote/DRjAT/src/lib/lib.jl:206 [inlined]
 [42] (::Zygote.var"#1912#back#210"{Zygote.var"#208#209"{Tuple{Tuple{Nothing, Nothing, Nothing}}, typeof(∂(contract))}})(Δ::ITensor)
    @ Zygote ~/.julia/packages/ZygoteRules/AIbCs/src/adjoint.jl:67
 [43] Pullback
    @ ~/.julia/packages/ITensors/PisGD/src/itensor.jl:2086 [inlined]
 [44] (::typeof(∂(#*#239)))(Δ::ITensor)
    @ Zygote ~/.julia/packages/Zygote/DRjAT/src/compiler/interface2.jl:0
 [45] (::Zygote.var"#208#209"{Tuple{Tuple{Nothing, Nothing}, Tuple{Nothing, Nothing, Nothing}}, typeof(∂(#*#239))})(Δ::ITensor)
    @ Zygote ~/.julia/packages/Zygote/DRjAT/src/lib/lib.jl:206
 [46] #1912#back
    @ ~/.julia/packages/ZygoteRules/AIbCs/src/adjoint.jl:67 [inlined]
 [47] Pullback
    @ ~/.julia/packages/ITensors/PisGD/src/itensor.jl:2086 [inlined]
 [48] (::typeof(∂(*)))(Δ::ITensor)
    @ Zygote ~/.julia/packages/Zygote/DRjAT/src/compiler/interface2.jl:0
 [49] Pullback
    @ ./In[40]:7 [inlined]
 [50] (::typeof(∂(f)))(Δ::Float64)
    @ Zygote ~/.julia/packages/Zygote/DRjAT/src/compiler/interface2.jl:0
 [51] Pullback
    @ ./In[42]:1 [inlined]
 [52] (::typeof(∂(#55)))(Δ::Float64)
    @ Zygote ~/.julia/packages/Zygote/DRjAT/src/compiler/interface2.jl:0
 [53] (::Zygote.var"#60#61"{typeof(∂(#55))})(Δ::Float64)
    @ Zygote ~/.julia/packages/Zygote/DRjAT/src/compiler/interface.jl:45
 [54] (::Zygote.var"#64#65"{var"#55#56"})(x::ITensor)
    @ Zygote ~/.julia/packages/Zygote/DRjAT/src/compiler/interface.jl:104
 [55] top-level scope
    @ In[45]:1
 [56] eval
    @ ./boot.jl:373 [inlined]

Version information

• Output from versioninfo():

julia> versioninfo()
Julia Version 1.7.0
Commit 3bf9d17731* (2021-11-30 12:12 UTC)
Platform Info:
  OS: Linux (x86_64-pc-linux-gnu)
  CPU: Intel(R) Xeon(R) Gold 6148 CPU @ 2.40GHz
  WORD_SIZE: 64
  LIBM: libopenlibm
  LLVM: libLLVM-12.0.1 (ORCJIT, skylake-avx512)
Environment:
  JULIA_NUM_THREADS = 128

• Output from using Pkg; Pkg.status("ITensors"):

julia> using Pkg; Pkg.status("ITensors")
Status `~/.julia/environments/v1.7/Project.toml`
  [9136182c] ITensors v0.3.19

[emstoudenmire]

Thanks for letting us know & this definitely looks like a bug.

Not that it is the ideal fix, but as a temporary workaround can you make an dense ITensor with all 1's down the diagonal and use that instead of a delta? (You can make a separate function to produce this tensor and set the 1 elements in a for loop.)

If you find after making that change that the speed of your forward pass doesn't change much (i.e. if this part of your code doesn't become a bottleneck after) then it should work and your code won't be slowed down by the replacement.

Thank you. I've set up a diagonal dense ITensor and the code runs well. Here is an example.

using ITensors
using ChainRulesCore
using Zygote

function getDelta(x, y)
    ind1, ind2 = inds(x)[1], inds(y)[1]
    d = ITensor(Float32, [ind1, ind2])
    l = min(NDTensors.dim(ind1), NDTensors.dim(ind2))
    for i in 1:l
        d[ind1 => i, ind2 => i] = 1.0f0
    end
    return d
end

function f(x, y)
    return (x * getDelta(x, y) * y)[]
end

@non_differentiable getDelta(::Any...)

idx_1 = Index(2)
idx_2 = Index(3)
tens_1 = randomITensor(Float32, idx_1)
tens_2 = randomITensor(Float32, idx_2)

l = x -> f(x, tens_2)
@assert l(tens_1) == f(tens_1, tens_2)

l'(tens_1)

[emstoudenmire]

Glad the workaround is successful. We will leave this open as a bug since it would be ideal for the AD system to be able to handle delta tensors of course, so thanks for pointing this out.

As a minor note, it could be faster (while still correct) to do

    for i in 1:l
        d[i, i] = 1.0f0
    end

inside your getDelta function, since the order of indices doesn't matter in that line of the code and it keeps ITensor from having to check or permute the indices each time.

[mtfishman]

Thanks for the nice report @aastrashko! I know other code using delta tensors in gradients work ok so it looks like this is hitting a corner case. Thanks for the suggestion for a workaround @emstoudenmire.

[mtfishman]

Turns out this bug is independent of the gradient code, and is just a bug in the Dense * Diag contraction code, i.e. this outputs the same error:

using ITensors
i = Index(3)
j = Index(2)
A = randomITensor(i)
B = delta(i, j)
A * B

[emstoudenmire]

Oh I see - good catch. And I see there’s a PR already. Really good to know the AD can already handle this case and it’s just a fixable forward pass bug.

On Tue, Sep 6, 2022 at 12:24 PM Matt Fishman ***@***.***> wrote: Turns out this bug is independent of the gradient code, and is just a bug in the Dense * Diag contraction code, i.e. this outputs the same error: using ITensors i = Index(3) j = Index(2) A = randomITensor(i) B = delta(i, j) A * B — Reply to this email directly, view it on GitHub <#969 (comment)>, or unsubscribe <https://github.com/notifications/unsubscribe-auth/AAHJ3RL4SNZCAOGHW2ZWDCTV45V5BANCNFSM6AAAAAAQDTHFN4> . You are receiving this because you were mentioned.Message ID: ***@***.***>

-- -=Miles Stoudenmire=- ***@***.*** ***@***.*** http://itensor.org/miles/