Fix SVD rrule when ΔSis a ZeroTangent

Jutho · Oct 17, 2023 · 680faf3 · 680faf3
1 parent e912eae
commit 680faf3
Showing 1 changed file with 1 addition and 1 deletion.
diff --git a/ext/TensorKitChainRulesCoreExt.jl b/ext/TensorKitChainRulesCoreExt.jl
@@ -178,7 +178,7 @@ function svd_rev(U::AbstractMatrix, S::AbstractMatrix, V::AbstractMatrix, ΔU,
     F = _invert_S²(S, tol)
     S⁻¹ = pinv(S; atol=tol)
 
-    term = Diagonal(diag(ΔS))
+    term = ΔS isa ZeroTangent ? ΔS : Diagonal(diag(ΔS))
 
     J = F .* (U' * ΔU)
     term += (J + J') * S