Don't use diagview(Vector)

Author: kshyatt

ext/MatrixAlgebraKitMooncakeExt/MatrixAlgebraKitMooncakeExt.jl

@@ -136,7 +136,7 @@ for (f!, f, f_full, pb, adj) in (
             copy!(D, diagview(DV[1]))
             V = DV[2]
             function $adj(::Mooncake.NoRData)
-                $pb(dA, A, (D, V), (dD, nothing))
+                $pb(dA, A, (Diagonal(D), V), (Diagonal(dD), nothing))
                 MatrixAlgebraKit.zero!(dD)
                 return Mooncake.NoRData(), Mooncake.NoRData(), Mooncake.NoRData(), Mooncake.NoRData()
             end
@@ -152,8 +152,8 @@ for (f!, f, f_full, pb, adj) in (
             output = diagview(DV[1])
             output_codual = Mooncake.CoDual(output, Mooncake.zero_tangent(output))
             function $adj(::Mooncake.NoRData)
-                D_dD = Mooncake.arrayify(D_dD)
-                $pb(dA, A, (D, V), (dD, nothing))
+                D, dD = Mooncake.arrayify(output_codual)
+                $pb(dA, A, (Diagonal(D), V), (Diagonal(dD), nothing))
                 MatrixAlgebraKit.zero!(dD)
                 return Mooncake.NoRData(), Mooncake.NoRData(), Mooncake.NoRData()
             end
@@ -181,7 +181,7 @@ for (f, pb, adj) in (
             function $adj(dy::Tuple{Mooncake.NoRData, Mooncake.NoRData, T}) where {T <: Real}
                 Dtrunc, Vtrunc, ϵ = Mooncake.primal(output_codual)
                 dDtrunc_, dVtrunc_, dϵ = Mooncake.tangent(output_codual)
-                abs(dϵ) > MatrixAlgebraKit.defaulttol(dϵ) && @warn "Pullback for $f does not yet support non-zero tangent for the truncation error"
+                abs(dy[3]) > MatrixAlgebraKit.defaulttol(dy[3]) && @warn "Pullback for $f does not yet support non-zero tangent for the truncation error"
                 D, dD = Mooncake.arrayify(Dtrunc, dDtrunc_)
                 V, dV = Mooncake.arrayify(Vtrunc, dVtrunc_)
                 $pb(dA, A, (D, V), (dD, dV))
@@ -275,7 +275,7 @@ function Mooncake.rrule!!(::CoDual{typeof(MatrixAlgebraKit.svd_vals!)}, A_dA::Co
     U, nS, Vᴴ = svd_compact(A, Mooncake.primal(alg_dalg))
     copy!(S, diagview(nS))
     function dsvd_vals_adjoint(::Mooncake.NoRData)
-        svd_pullback!(dA, A, (U, S, Vᴴ), (nothing, dS, nothing))
+        svd_pullback!(dA, A, (U, Diagonal(S), Vᴴ), (nothing, Diagonal(dS), nothing))
         MatrixAlgebraKit.zero!(dS)
         return Mooncake.NoRData(), Mooncake.NoRData(), Mooncake.NoRData(), Mooncake.NoRData()
     end
@@ -294,7 +294,7 @@ function Mooncake.rrule!!(::CoDual{typeof(MatrixAlgebraKit.svd_vals)}, A_dA::CoD
     S_codual = Mooncake.CoDual(diagview(S), Mooncake.fdata(Mooncake.zero_tangent(S)))
     function dsvd_vals_adjoint(::Mooncake.NoRData)
         S, dS = Mooncake.arrayify(S_codual)
-        svd_pullback!(dA, A, (U, S, Vᴴ), (nothing, dS, nothing))
+        svd_pullback!(dA, A, (U, Diagonal(S), Vᴴ), (nothing, Diagonal(dS), nothing))
         MatrixAlgebraKit.zero!(dS)
         return Mooncake.NoRData(), Mooncake.NoRData(), Mooncake.NoRData()
     end
@@ -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(dϵ) > MatrixAlgebraKit.defaulttol(dϵ) && @warn "Pullback for svd_trunc! does not yet support non-zero tangent for the truncation error"
+        abs(dy[3]) > MatrixAlgebraKit.defaulttol(dy[3]) && @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_)

src/common/view.jl

@@ -1,7 +1,6 @@
 # diagind: provided by LinearAlgebra.jl
 diagview(D::Diagonal) = D.diag
 diagview(D::AbstractMatrix) = view(D, diagind(D))
-diagview(D::AbstractVector) = D
 
 # triangularind
 function lowertriangularind(A::AbstractMatrix)