Add rrule for DiagonalTensorMap and include tests (#234)

* Add `rrule` for `DiagonalTensorMap(::AbstractTensorMap)`

* include tests

Author: lkdvos

ext/TensorKitChainRulesCoreExt/constructors.jl

@@ -47,6 +47,18 @@ function ChainRulesCore.rrule(::Type{<:DiagonalTensorMap}, data::DenseVector, ar
     return D, DiagonalTensorMap_pullback
 end
 
+function ChainRulesCore.rrule(::Type{DiagonalTensorMap}, t::AbstractTensorMap)
+    d = DiagonalTensorMap(t)
+    function DiagonalTensorMap_pullback(Δd_)
+        Δt = similar(t) # no projector needed
+        for (c, b) in blocks(unthunk(Δd_))
+            copy!(block(Δt, c), Diagonal(b))
+        end
+        return NoTangent(), Δt
+    end
+    return d, DiagonalTensorMap_pullback
+end
+
 function ChainRulesCore.rrule(::typeof(Base.getproperty), t::TensorMap, prop::Symbol)
     if prop === :data
         function getdata_pullback(Δdata)

test/ad.jl

@@ -196,6 +196,9 @@ Vlist = ((ℂ^2, (ℂ^3)', ℂ^3, ℂ^2, (ℂ^2)'),
             test_rrule(DiagonalTensorMap, D.data, D.domain)
             test_rrule(Base.getproperty, D, :data)
             test_rrule(Base.getproperty, D1, :data)
+
+            test_rrule(DiagonalTensorMap, rand!(T1))
+            test_rrule(DiagonalTensorMap, randn!(T))
         end
     end