Add tests rrules converters

Jutho · Jun 28, 2024 · c15b7ab · c15b7ab
1 parent 1148fe9
commit c15b7ab
Show file tree

Hide file tree

Showing 2 changed files with 9 additions and 3 deletions.
diff --git a/ext/TensorKitChainRulesCoreExt/constructors.jl b/ext/TensorKitChainRulesCoreExt/constructors.jl
@@ -6,8 +6,8 @@
 
 function ChainRulesCore.rrule(::Type{<:TensorMap}, d::DenseArray, args...)
     function TensorMap_pullback(Δt)
-        ∂d = convert(Array, Δt)
-        return NoTangent(), ∂d, fill(NoTangent(), length(args))...
+        ∂d = convert(Array, unthunk(Δt))
+        return NoTangent(), ∂d, ntuple(_ -> NoTangent(), length(args))...
     end
     return TensorMap(d, args...), TensorMap_pullback
 end
@@ -21,7 +21,7 @@ function ChainRulesCore.rrule(::typeof(Base.convert), T::Type{<:Array},
                               t::AbstractTensorMap)
     A = convert(T, t)
     function convert_pullback(ΔA)
-        ∂t = TensorMap(ΔA, codomain(t), domain(t))
+        ∂t = TensorMap(unthunk(ΔA), codomain(t), domain(t))
         return NoTangent(), NoTangent(), ∂t
     end
     return A, convert_pullback

diff --git a/test/ad.jl b/test/ad.jl
@@ -139,6 +139,12 @@ Vlist = ((ℂ^2, (ℂ^3)', ℂ^3, ℂ^2, (ℂ^2)'),
         P1 = ProjectTo(T1)
         @test P1(T1) == T1
         @test P1(T2) == real(T2)
+
+        test_rrule(copy, T1)
+        test_rrule(copy, T2)
+
+        test_rrule(convert, Array, T1)
+        test_rrule(TensorMap, convert(Array, T1), space(T1))
     end
 
     @timedtestset "Basic Linear Algebra with scalartype $T" for T in (Float64, ComplexF64)