Fix JET issues and Mooncake fails

Author: kshyatt

src/implementations/polar.jl

@@ -117,6 +117,7 @@ end
 function _left_polarnewton!(A::AbstractMatrix, W, P = similar(A, (0, 0)); tol = defaulttol(A), maxiter = 10)
     m, n = size(A) # we must have m >= n
     Rᴴinv = isempty(P) ? similar(P, (n, n)) : P # use P as workspace when available
+    Q = similar(A, (0, 0))
     if m > n # initial QR
         Q, R = qr_compact!(A)
         Rc = view(A, 1:n, 1:n)
@@ -157,6 +158,7 @@ end
 function _right_polarnewton!(A::AbstractMatrix, Wᴴ, P = similar(A, (0, 0)); tol = defaulttol(A), maxiter = 10)
     m, n = size(A) # we must have m <= n
     Lᴴinv = isempty(P) ? similar(P, (m, m)) : P # use P as workspace when available
+    Q = similar(A, (0, 0))
     if m < n # initial QR
         L, Q = lq_compact!(A)
         Lc = view(A, 1:m, 1:m)

src/implementations/qr.jl

@@ -144,6 +144,7 @@ function _lapack_qr!(
         throw(ArgumentError("inplace Q only supported if matrix is tall (`m >= n`), R is not required, and using the unblocked algorithm (`blocksize=1`) with `positive=false`"))
     end
 
+    jpvt = Vector{Int}(undef, 0)
     if blocksize > 1
         nb = min(minmn, blocksize)
         if computeR # first use R as space for T

src/pullbacks/eig.jl

@@ -125,7 +125,7 @@ function eig_trunc_pullback!(
     p == length(D) || throw(DimensionMismatch())
     (n, n) == size(ΔA) || throw(DimensionMismatch())
     G = V' * V
-
+    ΔVperp = similar(V, (0, 0))
     if !iszerotangent(ΔV)
         (n, p) == size(ΔV) || throw(DimensionMismatch())
         VᴴΔV = V' * ΔV

src/pullbacks/svd.jl

@@ -42,6 +42,8 @@ function svd_pullback!(
     ΔU, ΔSmat, ΔVᴴ = ΔUSVᴴ
     UΔU = fill!(similar(U, (r, r)), 0)
     VΔV = fill!(similar(Vᴴ, (r, r)), 0)
+    indU = axes(U, 2)
+    indV = axes(Vᴴ, 1)
     if !iszerotangent(ΔU)
         m == size(ΔU, 1) || throw(DimensionMismatch("first dimension of ΔU ($(size(ΔU, 1))) does not match first dimension of U ($m)"))
         pU = size(ΔU, 2)

test/mooncake.jl

@@ -9,13 +9,9 @@ using LinearAlgebra: UpperTriangular, Diagonal, Hermitian, mul!
 
 include("ad_utils.jl")
 
-make_mooncake_tangent(ΔAelem::T) where {T <: Complex} = Mooncake.build_tangent(T, real(ΔAelem), imag(ΔAelem))
-make_mooncake_tangent(ΔA::Matrix{<:Real}) = ΔA
-make_mooncake_tangent(ΔA::Vector{<:Real}) = ΔA
-make_mooncake_tangent(ΔA::Matrix{T}) where {T <: Complex} = map(make_mooncake_tangent, ΔA)
-make_mooncake_tangent(ΔA::Vector{T}) where {T <: Complex} = map(make_mooncake_tangent, ΔA)
-make_mooncake_tangent(ΔD::Diagonal{T}) where {T <: Real} = Mooncake.build_tangent(typeof(ΔD), diagview(ΔD))
-make_mooncake_tangent(ΔD::Diagonal{T}) where {T <: Complex} = Mooncake.build_tangent(typeof(ΔD), map(make_mooncake_tangent, diagview(ΔD)))
+make_mooncake_tangent(ΔA::Matrix) = ΔA
+make_mooncake_tangent(ΔA::Vector) = ΔA
+make_mooncake_tangent(ΔD::Diagonal) = Mooncake.build_tangent(typeof(ΔD), diagview(ΔD))
 
 make_mooncake_tangent(T::Tuple) = Mooncake.build_tangent(typeof(T), T...)