MatrixAlgebraKit.jl/src/pullbacks/polar.jl at 141b1fa99760b0b2cab97c5def00cb6ef41e75ec · QuantumKitHub/MatrixAlgebraKit.jl · GitHub

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
"""
    left_polar_pullback!(ΔA, A, WP, ΔWP)

Adds the pullback from the left polar decomposition of `A` to `ΔA` given the output `WP` and
cotangent `ΔWP` of `left_polar(A)`.
"""
function left_polar_pullback!(ΔA::AbstractMatrix, A, WP, ΔWP; kwargs...)
    # Extract the Polar components
    W, P = WP

    # Extract and check the cotangents
    ΔW, ΔP = ΔWP
    if !iszerotangent(ΔP)
        ΔP = project_hermitian(ΔP)
    end
    M = zero(P)
    !iszerotangent(ΔW) && mul!(M, W', ΔW, 1, 1)
    !iszerotangent(ΔP) && mul!(M, ΔP, P, -1, 1)
    C = sylvester(P, P, M' - M)
    C .+= ΔP
    ΔA = mul!(ΔA, W, C, 1, 1)
    if !iszerotangent(ΔW)
        ΔWP = ΔW / P
        WdΔWP = W' * ΔWP
        ΔWP = mul!(ΔWP, W, WdΔWP, -1, 1)
        ΔA .+= ΔWP
    end
    return ΔA
end

"""
    right_polar_pullback!(ΔA, A, PWᴴ, ΔPWᴴ)

Adds the pullback from the left polar decomposition of `A` to `ΔA` given the output `PWᴴ`
and cotangent `ΔPWᴴ` of `right_polar(A)`.
"""
function right_polar_pullback!(ΔA::AbstractMatrix, A, PWᴴ, ΔPWᴴ; kwargs...)
    # Extract the Polar components
    P, Wᴴ = PWᴴ

    # Extract and check the cotangents
    ΔP, ΔWᴴ = ΔPWᴴ
    if !iszerotangent(ΔP)
        ΔP = project_hermitian(ΔP)
    end
    M = zero(P)
    !iszerotangent(ΔWᴴ) && mul!(M, ΔWᴴ, Wᴴ', 1, 1)
    !iszerotangent(ΔP) && mul!(M, P, ΔP, -1, 1)
    C = sylvester(P, P, M' - M)
    C .+= ΔP
    ΔA = mul!(ΔA, C, Wᴴ, 1, 1)
    if !iszerotangent(ΔWᴴ)
        PΔWᴴ = P \ ΔWᴴ
        PΔWᴴW = PΔWᴴ * Wᴴ'
        PΔWᴴ = mul!(PΔWᴴ, PΔWᴴW, Wᴴ, -1, 1)
        ΔA .+= PΔWᴴ
    end
    return ΔA
end