Tangent type and bad index

Author: kshyatt

ext/MatrixAlgebraKitMooncakeExt/MatrixAlgebraKitMooncakeExt.jl

@@ -291,7 +291,7 @@ function Mooncake.rrule!!(::CoDual{typeof(MatrixAlgebraKit.svd_vals)}, A_dA::CoD
     # of ComplexF32) into the correct **forwards** data type (since we are now in the forward
     # pass). For many types this is done automatically when the forward step returns, but
     # not for nested structs with various fields (like Diagonal{Complex})
-    S_codual = Mooncake.CoDual(diagview(S), Mooncake.fdata(Mooncake.zero_tangent(S)))
+    S_codual = Mooncake.CoDual(diagview(S), Mooncake.fdata(Mooncake.zero_tangent(diagview(S))))
     function dsvd_vals_adjoint(::Mooncake.NoRData)
         S, dS = Mooncake.arrayify(S_codual)
         svd_pullback!(dA, A, (U, Diagonal(S), Vᴴ), (nothing, Diagonal(dS), nothing))
@@ -317,7 +317,7 @@ function Mooncake.rrule!!(::CoDual{typeof(MatrixAlgebraKit.svd_trunc)}, A_dA::Co
     function dsvd_trunc_adjoint(dy::Tuple{Mooncake.NoRData, Mooncake.NoRData, Mooncake.NoRData, T}) where {T <: Real}
         Utrunc, Strunc, Vᴴtrunc, ϵ = Mooncake.primal(output_codual)
         dUtrunc_, dStrunc_, dVᴴtrunc_, dϵ = Mooncake.tangent(output_codual)
-        abs(dy[3]) > MatrixAlgebraKit.defaulttol(dy[3]) && @warn "Pullback for svd_trunc! does not yet support non-zero tangent for the truncation error"
+        abs(dy[4]) > MatrixAlgebraKit.defaulttol(dy[4]) && @warn "Pullback for svd_trunc! does not yet support non-zero tangent for the truncation error"
         U, dU = Mooncake.arrayify(Utrunc, dUtrunc_)
         S, dS = Mooncake.arrayify(Strunc, dStrunc_)
         Vᴴ, dVᴴ = Mooncake.arrayify(Vᴴtrunc, dVᴴtrunc_)