update qr/lq_null pullbacks

Author: Jutho

ext/MatrixAlgebraKitChainRulesCoreExt.jl

@@ -43,11 +43,10 @@ for qr_f in (:qr_compact, :qr_full)
 end
 function ChainRulesCore.rrule(::typeof(qr_null!), A::AbstractMatrix, N, alg)
     Ac = copy_input(qr_full, A)
-    QR = qr_full!(Ac, initialize_output(qr_full!, A, alg), alg)
-    N = copy!(N, view(QR[1], 1:size(A, 1), (size(A, 2) + 1):size(A, 1)))
+    N = qr_null!(Ac, N, alg)
     function qr_null_pullback(ΔN)
         ΔA = zero(A)
-        MatrixAlgebraKit.qr_null_pullback!(ΔA, A, QR, unthunk(ΔN))
+        MatrixAlgebraKit.qr_null_pullback!(ΔA, A, N, unthunk(ΔN))
         return NoTangent(), ΔA, ZeroTangent(), NoTangent()
     end
     function qr_null_pullback(::ZeroTangent) # is this extra definition useful?
@@ -76,11 +75,10 @@ for lq_f in (:lq_compact, :lq_full)
 end
 function ChainRulesCore.rrule(::typeof(lq_null!), A::AbstractMatrix, Nᴴ, alg)
     Ac = copy_input(lq_full, A)
-    LQ = lq_full!(Ac, initialize_output(lq_full!, A, alg), alg)
-    Nᴴ = copy!(Nᴴ, view(LQ[2], (size(A, 1) + 1):size(A, 2), 1:size(A, 2)))
+    Nᴴ = lq_null!(Ac, Nᴴ, alg)
     function lq_null_pullback(ΔNᴴ)
         ΔA = zero(A)
-        MatrixAlgebraKit.lq_null_pullback!(ΔA, A, LQ, unthunk(ΔNᴴ))
+        MatrixAlgebraKit.lq_null_pullback!(ΔA, A, Nᴴ, unthunk(ΔNᴴ))
         return NoTangent(), ΔA, ZeroTangent(), NoTangent()
     end
     function lq_null_pullback(::ZeroTangent) # is this extra definition useful?

src/pullbacks/lq.jl

@@ -105,18 +105,26 @@ function lq_compact_pullback!(
 end
 
 """
-    lq_null_pullback(ΔA, A, LQ, ΔNᴴ)
+    lq_null_pullback(ΔA, A, Nᴴ, ΔNᴴ)
 
-Adds the pullback from the nullspace of the LQ decomposition of `A` to `ΔA` given the
-factorization `LQ` and cotangent `ΔNᴴ` of `lq_null(A)`.
+Adds the pullback from the left nullspace of `A` to `ΔA`, given the nullspace basis
+ `Nᴴ` and its cotangent `ΔNᴴ` of `lq_null(A)`.
 
 See also [`lq_compact_pullback!`](@ref).
 """
-function lq_null_pullback!(ΔA::AbstractMatrix, A, LQ, ΔNᴴ)
-    (m, n) = size(A)
-    minmn = min(m, n)
-    ΔQ = zero!(similar(A, (n, n)))
-    view(ΔQ, (minmn + 1):n, 1:n) .= ΔNᴴ
-    lq_compact_pullback!(ΔA, A, LQ, (nothing, ΔQ))
+function lq_null_pullback!(
+        ΔA::AbstractMatrix, A, Nᴴ, ΔNᴴ;
+        tol::Real = default_pullback_gaugetol(A),
+        gauge_atol::Real = tol
+    )
+    if !iszerotangent(ΔNᴴ) && size(Nᴴ, 1) > 0
+        NᴴΔN = Nᴴ * ΔNᴴ'
+        Δgauge = norm((NᴴΔN .- NᴴΔN') ./ 2)
+        Δgauge < tol ||
+            @warn "`lq_null` cotangent sensitive to gauge choice: (|Δgauge| = $Δgauge)"
+        L, Q = lq_compact(A; positive = true) # should we be able to provide algorithm here?
+        X = ldiv!(LowerTriangular(L)', Q * ΔNᴴ')
+        ΔA = mul!(ΔA, X, Nᴴ, -1, 1)
+    end
     return ΔA
 end

src/pullbacks/qr.jl

@@ -104,17 +104,27 @@ function qr_compact_pullback!(
 end
 
 """
-    qr_null_pullback(ΔA, A, QR, ΔNᴴ)
+    qr_null_pullback(ΔA, A, N, ΔN)
 
-Adds the pullback from the nullspace of the QR decomposition of `A` to `ΔA` given the
-factorization `QR` and cotangent `ΔN` of `qr_null(A)`.
+Adds the pullback from the right nullspace of `A` to `ΔA`, given the nullspace basis
+`N` and its cotangent `ΔN` of `qr_null(A)`.
 
 See also [`qr_compact_pullback!`](@ref).
 """
-function qr_null_pullback!(ΔA::AbstractMatrix, A, QR, ΔN)
-    m, n = size(A)
-    minmn = min(m, n)
-    ΔQ = zero!(similar(A, (m, m)))
-    view(ΔQ, 1:m, (minmn + 1):m) .= ΔN
-    return qr_compact_pullback!(ΔA, A, QR, (ΔQ, nothing))
+function qr_null_pullback!(
+        ΔA::AbstractMatrix, A, N, ΔN;
+        tol::Real = default_pullback_gaugetol(A),
+        gauge_atol::Real = tol
+    )
+    if !iszerotangent(ΔN) && size(N, 2) > 0
+        NᴴΔN = N' * ΔN
+        Δgauge = norm((NᴴΔN .- NᴴΔN') ./ 2)
+        Δgauge < tol ||
+            @warn "`qr_null` cotangent sensitive to gauge choice: (|Δgauge| = $Δgauge)"
+
+        Q, R = qr_compact(A; positive = true)
+        X = rdiv!(ΔN' * Q, UpperTriangular(R)')
+        ΔA = mul!(ΔA, N, X, -1, 1)
+    end
+    return ΔA
 end