Less lazy ProjectTo fix

Jutho · Oct 16, 2023 · 484a323 · 484a323
1 parent df0091f
commit 484a323
Showing 1 changed file with 10 additions and 1 deletion.
diff --git a/ext/TensorKitChainRulesCoreExt.jl b/ext/TensorKitChainRulesCoreExt.jl
@@ -59,7 +59,16 @@ function ChainRulesCore.rrule(::typeof(Base.copy), t::AbstractTensorMap)
 end
 
 ChainRulesCore.ProjectTo(::T) where {T<:AbstractTensorMap} = ProjectTo{T}()
-(::ProjectTo{T})(x::AbstractTensorMap) where {T<:AbstractTensorMap} = convert(T, x)
+function (::ProjectTo{T1})(x::T2) where {S,N1,N2,T1<:AbstractTensorMap{S,N1,N2},
+                                         T2<:AbstractTensorMap{S,N1,N2}}
+    T1 === T2 && return x
+    y = similar(x, scalartype(T1))
+    for (c, b) in blocks(y)
+        p = ProjectTo(b)
+        b .= p(block(x, c))
+    end
+    return y
+end
 
 # Base Linear Algebra
 # -------------------