Fix rrules of TensorOperations with DiagonalTensorMap

[lkdvos]

Rewrites these rules in terms of similar instead of zerovector to ensure first contracting, then projecting onto a diagonal input.

Fixes #209

[codecov]

Codecov Report

Attention: Patch coverage is 90.90909% with 2 lines in your changes missing coverage. Please review.

Project coverage is 82.25%. Comparing base (a264735) to head (5fd86c0).
Report is 1 commits behind head on master.

Files with missing lines	Patch %	Lines
ext/TensorKitChainRulesCoreExt/utility.jl	60.00%	2 Missing ⚠️

Additional details and impacted files

@@            Coverage Diff             @@
##           master     #210      +/-   ##
==========================================
+ Coverage   82.17%   82.25%   +0.08%     
==========================================
  Files          43       43              
  Lines        5424     5433       +9     
==========================================
+ Hits         4457     4469      +12     
+ Misses        967      964       -3     

☔ View full report in Codecov by Sentry.
📢 Have feedback on the report? Share it here.

[Jutho]

Looks good to me, but I wonder about the similar semantics. In LinearAlgebra.jl, similar(d::Diagonal) produces a new Diagonal.

[lkdvos]

Honestly that function is such a mess... From the docstring, it should return a mutable array of the given eltype and dimensions, which is precisely what I want here. LinearAlgebra strikes again with their inconsistency:

julia> d = Diagonal(rand(2));

julia> similar(d)
2×2 Diagonal{Float64, Vector{Float64}}:
 5.0e-324   ⋅
  ⋅        1.0e-323

julia> similar(d, size(d))
2×2 Matrix{Float64}:
   0.0  5.0e-324
 NaN    0.0

julia> similar(d, Float32)
2×2 Diagonal{Float32, Vector{Float32}}:
 -1.33821f-27   ⋅
   ⋅           4.0f-45

julia> similar(d, Float32, size(d))
2×2 Matrix{Float32}:
 2.11663f-24  2.20105f-23
 1.0f-45      1.0f-45

In principle, I could manually hook into the tensoralloc machinery as well, but I don't know if that's really necessary here?

[Jutho]

So even if we choose to mimick the similar(::Diagonal, ...) behaviour (not saying that we must, just because some inconsistencies between LinearAlgebra and TensorKit were pointed out to me in the past), the following approach (i.e. adding an explicit space argument) would ensure that we don't get a DiagonalTensor:

_dA = similar(A, promote_contract(scalartype(ΔC), scalartype(B), scalartype(α)), space(A))

[lkdvos]

So I cracked and just wrote everything in terms of tensoralloc, I'm assuming that to be the most stable in the long run, no matter what we decide on for similar.

I did find that we silently ignore the partition of the pAB argument for the in-place tensor operations, and just take the partition as imposed by C. I'm not sure this is the right choice, and we might to explicitly error, but I'm okay to leave this as is for now.

[Jutho]

Yes, that's a choice I made a long time ago; to facilitate the use case where a user doesn't want to think about a tensor having a bipartition into two sets of indices and thus uses @tensor with a single list of indices for each tensor, but then still uses TensorMap objects where both domain and codomain are both nontrivial, because they come for example from a tsvd call. So pAB is respected when a new output tensor is created, but not in the inplace version. For the planar methods, that should be different though and there an error should be thrown.

[Jutho]

Ok so it seems that there is some type instability in the local variables of the anonymous functions that go into the thunk of dA and dB for tensorcontract!. This is the distilled output (true vs inferred) for dA:

Thunk{TensorKitChainRulesCoreExt.var"#90#105"{Tuple{Tuple{Int64, Int64}, Tuple{Int64, Int64}}, TensorMap{ComplexF64, ComplexSpace, 3, 1, Vector{ComplexF64}}, Tensor{ComplexF64, ComplexSpace, 3, Vector{ComplexF64}}, Tuple{Tuple{Int64, Int64}, Tuple{Int64}}, Bool, TensorMap{ComplexF64, ComplexSpace, 2, 1, Vector{ComplexF64}}, Tuple{Tuple{Int64}, Tuple{Int64, Int64}}, Bool, ComplexF64, Tuple{}, ProjectTo{Tensor{ComplexF64, ComplexSpace, 3, Vector{ComplexF64}}, @NamedTuple{}}}}
Thunk{F} where F<:(TensorKitChainRulesCoreExt.var"#90#105"{<:Tuple{Union{Tuple{}, Tuple{Int64, Vararg{Int64}}}, Union{Tuple{}, Tuple{Int64, Vararg{Int64}}}}, TensorMap{ComplexF64, ComplexSpace, 3, 1, Vector{ComplexF64}}, Tensor{ComplexF64, ComplexSpace, 3, Vector{ComplexF64}}, Tuple{Tuple{Int64, Int64}, Tuple{Int64}}, Bool, TensorMap{ComplexF64, ComplexSpace, 2, 1, Vector{ComplexF64}}, Tuple{Tuple{Int64}, Tuple{Int64, Int64}}, Bool, ComplexF64, Tuple{}, ProjectTo{Tensor{ComplexF64, ComplexSpace, 3, Vector{ComplexF64}}, @NamedTuple{}}})

and for dB

Thunk{TensorKitChainRulesCoreExt.var"#95#110"{Tuple{Tuple{Int64, Int64}, Tuple{Int64, Int64}}, TensorMap{ComplexF64, ComplexSpace, 3, 1, Vector{ComplexF64}}, Tensor{ComplexF64, ComplexSpace, 3, Vector{ComplexF64}}, Tuple{Tuple{Int64, Int64}, Tuple{Int64}}, Bool, TensorMap{ComplexF64, ComplexSpace, 2, 1, Vector{ComplexF64}}, Tuple{Tuple{Int64}, Tuple{Int64, Int64}}, Bool, ComplexF64, Tuple{}, ProjectTo{TensorMap{ComplexF64, ComplexSpace, 2, 1, Vector{ComplexF64}}, @NamedTuple{}}}}
Thunk{F} where F<:(TensorKitChainRulesCoreExt.var"#95#110"{<:Tuple{Union{Tuple{}, Tuple{Int64, Vararg{Int64}}}, Union{Tuple{}, Tuple{Int64, Vararg{Int64}}}}, TensorMap{ComplexF64, ComplexSpace, 3, 1, Vector{ComplexF64}}, Tensor{ComplexF64, ComplexSpace, 3, Vector{ComplexF64}}, Tuple{Tuple{Int64, Int64}, Tuple{Int64}}, Bool, TensorMap{ComplexF64, ComplexSpace, 2, 1, Vector{ComplexF64}}, Tuple{Tuple{Int64}, Tuple{Int64, Int64}}, Bool, ComplexF64, Tuple{}, ProjectTo{TensorMap{ComplexF64, ComplexSpace, 2, 1, Vector{ComplexF64}}, @NamedTuple{}}}),

[Jutho]

I hope my latest commit fixes this. I don't know why now all of a sudden the @constprop :aggressive statements are necessary. It seems we are turning into an aggressive world 😄 .

[lkdvos]

Thanks for fixing this in any case. I guess that function should have been using Val, but it's really unfortunate the compiler doesn't realize that. I guess @inline could have also done the trick, but I'm definitely okay with this.

[Jutho]

I am also fine what any of those suggestions, whatever you think is most future-proof (probably Val). Anyway, for me this is good to merge or to change to Val and then merge whenever you want.