Format

Author: kshyatt

ext/MatrixAlgebraKitMooncakeExt/MatrixAlgebraKitMooncakeExt.jl

@@ -76,8 +76,8 @@ for (f, f_full, pb, adj) in (
             arg = $f(copy(A), arg, Mooncake.primal(alg_dalg))
             function $adj(::Mooncake.NoRData)
                 $pb(dA, A, arg, darg; kwargs...)
-                A    .= Ac
-                arg  .= argc
+                A .= Ac
+                arg .= argc
                 MatrixAlgebraKit.zero!(darg)
                 return Mooncake.NoRData(), Mooncake.NoRData(), Mooncake.NoRData(), Mooncake.NoRData()
             end
@@ -89,17 +89,17 @@ end
 @is_primitive Mooncake.DefaultCtx Mooncake.ReverseMode Tuple{typeof(MatrixAlgebraKit.eig_vals!), Any, Any, MatrixAlgebraKit.AbstractAlgorithm}
 function Mooncake.rrule!!(::CoDual{<:typeof(MatrixAlgebraKit.eig_vals!)}, A_dA::CoDual, D_dD::CoDual, alg_dalg::CoDual; kwargs...)
     # compute primal
-    D_    = Mooncake.primal(D_dD)
-    dD_   = Mooncake.tangent(D_dD)
-    A_    = Mooncake.primal(A_dA)
-    dA_   = Mooncake.tangent(A_dA)
+    D_ = Mooncake.primal(D_dD)
+    dD_ = Mooncake.tangent(D_dD)
+    A_ = Mooncake.primal(A_dA)
+    dA_ = Mooncake.tangent(A_dA)
     A, dA = arrayify(A_, dA_)
     D, dD = arrayify(D_, dD_)
-    Ac    = copy(A)
-    Dc    = copy(D)
+    Ac = copy(A)
+    Dc = copy(D)
     # update primal
-    DV    = eig_full(A, Mooncake.primal(alg_dalg); kwargs...)
-    V     = DV[2]
+    DV = eig_full(A, Mooncake.primal(alg_dalg); kwargs...)
+    V = DV[2]
     eig_vals!(A, D, Mooncake.primal(alg_dalg))
     function deig_vals_adjoint(::Mooncake.NoRData)
         A .= Ac
@@ -114,19 +114,19 @@ end
 @is_primitive Mooncake.DefaultCtx Mooncake.ReverseMode Tuple{typeof(MatrixAlgebraKit.eigh_vals!), Any, Any, MatrixAlgebraKit.AbstractAlgorithm}
 function Mooncake.rrule!!(::CoDual{<:typeof(MatrixAlgebraKit.eigh_vals!)}, A_dA::CoDual, D_dD::CoDual, alg_dalg::CoDual; kwargs...)
     # compute primal
-    D_    = Mooncake.primal(D_dD)
-    dD_   = Mooncake.tangent(D_dD)
-    A_    = Mooncake.primal(A_dA)
-    dA_   = Mooncake.tangent(A_dA)
+    D_ = Mooncake.primal(D_dD)
+    dD_ = Mooncake.tangent(D_dD)
+    A_ = Mooncake.primal(A_dA)
+    dA_ = Mooncake.tangent(A_dA)
     A, dA = arrayify(A_, dA_)
-    Ac    = copy(A)
+    Ac = copy(A)
     D, dD = arrayify(D_, dD_)
-    Dc    = copy(D)
-    DV    = eigh_full(A, Mooncake.primal(alg_dalg); kwargs...)
-    D    .= diagview(DV[1])
-    V     = DV[2]
+    Dc = copy(D)
+    DV = eigh_full(A, Mooncake.primal(alg_dalg); kwargs...)
+    D .= diagview(DV[1])
+    V = DV[2]
     function deigh_vals_adjoint(::Mooncake.NoRData)
-        A  .= Ac
+        A .= Ac
         eigh_pullback!(dA, A, (D, V), (dD, nothing); kwargs...)
         D .= Dc
         MatrixAlgebraKit.zero!(dD)
@@ -154,7 +154,7 @@ for f in (svd_full!, svd_compact!)
             minmn = min(size(A)...)
             function dsvd_adjoint(::Mooncake.NoRData)
                 A .= Ac
-                if ($f == svd_compact!) 
+                if ($f == svd_compact!)
                     svd_pullback!(dA, A, (U, S, Vᴴ), (dU, dS, dVᴴ))
                 else # full
                     vU = view(U, :, 1:minmn)
@@ -165,8 +165,8 @@ for f in (svd_full!, svd_compact!)
                     vdVᴴ = view(dVᴴ, 1:minmn, :)
                     svd_pullback!(dA, A, (vU, vS, vVᴴ), (vdU, vdS, vdVᴴ))
                 end
-                U  .= Uc
-                S  .= Sc
+                U .= Uc
+                S .= Sc
                 Vᴴ .= Vᴴc
                 MatrixAlgebraKit.zero!(dU)
                 MatrixAlgebraKit.zero!(dS)
@@ -181,17 +181,17 @@ end
 @is_primitive Mooncake.DefaultCtx Mooncake.ReverseMode Tuple{typeof(MatrixAlgebraKit.svd_vals!), Any, Any, MatrixAlgebraKit.AbstractAlgorithm}
 function Mooncake.rrule!!(::CoDual{<:typeof(MatrixAlgebraKit.svd_vals!)}, A_dA::CoDual, S_dS::CoDual, alg_dalg::CoDual; kwargs...)
     # compute primal
-    S_  = Mooncake.primal(S_dS)
+    S_ = Mooncake.primal(S_dS)
     dS_ = Mooncake.tangent(S_dS)
-    A_  = Mooncake.primal(A_dA)
+    A_ = Mooncake.primal(A_dA)
     dA_ = Mooncake.tangent(A_dA)
     A, dA = arrayify(A_, dA_)
     S, dS = arrayify(S_, dS_)
-    Ac    = copy(A)
+    Ac = copy(A)
     U, nS, Vᴴ = svd_compact(A, Mooncake.primal(alg_dalg); kwargs...)
     S .= diagview(nS)
     function dsvd_vals_adjoint(::Mooncake.NoRData)
-        A  .= Ac
+        A .= Ac
         svd_pullback!(dA, A, (U, S, Vᴴ), (nothing, dS, nothing))
         MatrixAlgebraKit.zero!(dS)
         return Mooncake.NoRData(), Mooncake.NoRData(), Mooncake.NoRData(), Mooncake.NoRData()

src/pullbacks/eig.jl

@@ -47,7 +47,7 @@ function eig_pullback!(
         Δgauge = norm(view(VᴴΔV, mask), Inf)
         Δgauge < gauge_atol ||
             @warn "`eig` cotangents sensitive to gauge choice: (|Δgauge| = $Δgauge)"
-        
+
         VᴴΔV .*= conj.(inv_safe.(transpose(D) .- D, tol))
 
         if !iszerotangent(ΔDmat)