QR with pivoting? about backwardslinalg.jl

[GiggleLiu]

For non-square, complex valued matrix, yes, but it does not support pivoting yet.

For pivoted qr, the computational graph looks like

A -> * P' -> qr -> Q, R

Notice P matrix is just a permutation matrix, it does not contain anay gradient information. So the backward rule would be simplely dQ, dR -> back(qr) -> * P -> dA.

But it would be nice to have an optional parameter. Thanks for your issue!

[GiggleLiu]

e.g. define the adjoint like this (will be merged soon)

using BackwardsLinalg: qr_back, copyltu!, gradient_check
using Zygote
using Zygote: @adjoint
@adjoint function qr(A::AbstractMatrix, pivot::Val{true})
    res = qr(A, pivot)
    Q, R, P = Matrix(res.Q), res.R, res.P
    @show P
    (Q, R, P), dy -> (qr_back(Q*R, Q, R, dy[1], dy[2])*P',nothing)
end

using Random, Test
@testset "qr Q complex pivot" begin
    T = Float64
    Random.seed!(3)
    for (M, N) in [(20, 60), (40, 40), (60, 20)]
        A = randn(T, M, N)
        op = randn(M, M)
        op += op'
        op2 = randn(N, N)
        op2 += op2'

        function tfunc(x)
            Q, R, P = qr(x, Val(true))
            v = Q[:,1]
            v2 = R[2,:]
            @show size(v), size(v2), size(op), size(op2)
            (v'*op*v + v2'*op2*v2)[] |> real
        end
        @show tfunc(A)
        @test gradient_check(tfunc, A)
    end
    a = [1+1im 2+1im; 3-1im 4+2im]
    @test copyltu!(a) ≈ [1 3+1im; 3-1im 4]
end
``

[GiggleLiu]

I have pushed to master directly (since no one is reviewing me)
see this commit for reference.

Note: all these functions are pirating LinearAlgebra, I will try to avoid this in 0.2.