diff --git a/ext/TensorKitChainRulesCoreExt/constructors.jl b/ext/TensorKitChainRulesCoreExt/constructors.jl index 81cd2dbe9..b6ef760da 100644 --- a/ext/TensorKitChainRulesCoreExt/constructors.jl +++ b/ext/TensorKitChainRulesCoreExt/constructors.jl @@ -57,6 +57,13 @@ function ChainRulesCore.rrule(::Type{DiagonalTensorMap}, t::AbstractTensorMap) return d, DiagonalTensorMap_pullback end +function ChainRulesCore.rrule(::Type{TensorMap}, d::DiagonalTensorMap) + t = TensorMap(d) + P = ProjectTo(d) + TensorMap_pullback(Δt) = NoTangent(), P(unthunk(Δt)) + return t, TensorMap_pullback +end + function ChainRulesCore.rrule(::typeof(Base.getproperty), t::TensorMap, prop::Symbol) if prop === :data function getdata_pullback(Δdata) diff --git a/test/chainrules/linalg.jl b/test/chainrules/linalg.jl index cd3da8b2a..880412a3c 100644 --- a/test/chainrules/linalg.jl +++ b/test/chainrules/linalg.jl @@ -89,6 +89,9 @@ for V in spacelist test_rrule(DiagonalTensorMap, rand!(T1)) test_rrule(DiagonalTensorMap, randn!(T)) + + test_rrule(TensorMap, D1) + test_rrule(TensorMap, D) end end