Automatic Differentiation

Project.toml

@@ -7,12 +7,19 @@ version = "0.11.2"
 HalfIntegers = "f0d1745a-41c9-11e9-1dd9-e5d34d218721"
 LRUCache = "8ac3fa9e-de4c-5943-b1dc-09c6b5f20637"
 LinearAlgebra = "37e2e46d-f89d-539d-b4ee-838fcccc9c8e"
+PackageExtensionCompat = "65ce6f38-6b18-4e1d-a461-8949797d7930"
 Strided = "5e0ebb24-38b0-5f93-81fe-25c709ecae67"
 TensorOperations = "6aa20fa7-93e2-5fca-9bc0-fbd0db3c71a2"
 TupleTools = "9d95972d-f1c8-5527-a6e0-b4b365fa01f6"
 VectorInterface = "409d34a3-91d5-4945-b6ec-7529ddf182d8"
 WignerSymbols = "9f57e263-0b3d-5e2e-b1be-24f2bb48858b"
 
+[weakdeps]
+ChainRulesCore = "d360d2e6-b24c-11e9-a2a3-2a2ae2dbcce4"
+
+[extensions]
+TensorKitChainRulesCoreExt = "ChainRulesCore"
+
 [compat]
 HalfIntegers = "1"
 LRUCache = "1.0.2"
@@ -24,7 +31,10 @@ WignerSymbols = "1,2"
 julia = "1.6"
 
 [extras]
+ChainRulesCore = "d360d2e6-b24c-11e9-a2a3-2a2ae2dbcce4"
+ChainRulesTestUtils = "cdddcdb0-9152-4a09-a978-84456f9df70a"
 Combinatorics = "861a8166-3701-5b0c-9a16-15d98fcdc6aa"
+FiniteDifferences = "26cc04aa-876d-5657-8c51-4c34ba976000"
 HalfIntegers = "f0d1745a-41c9-11e9-1dd9-e5d34d218721"
 LinearAlgebra = "37e2e46d-f89d-539d-b4ee-838fcccc9c8e"
 Random = "9a3f8284-a2c9-5f02-9a11-845980a1fd5c"
@@ -34,4 +44,4 @@ TestExtras = "5ed8adda-3752-4e41-b88a-e8b09835ee3a"
 WignerSymbols = "9f57e263-0b3d-5e2e-b1be-24f2bb48858b"
 
 [targets]
-test = ["Combinatorics", "HalfIntegers", "LinearAlgebra", "Random", "TensorOperations", "Test", "TestExtras", "WignerSymbols"]
+test = ["Combinatorics", "HalfIntegers", "LinearAlgebra", "Random", "TensorOperations", "Test", "TestExtras", "WignerSymbols", "ChainRulesCore", "ChainRulesTestUtils", "FiniteDifferences"]

ext/TensorKitChainRulesCoreExt.jl

@@ -0,0 +1,337 @@
+module TensorKitChainRulesCoreExt
+
+using TensorOperations
+using TensorKit
+using ChainRulesCore
+using LinearAlgebra
+using TupleTools
+
+# Utility
+# -------
+
+_conj(conjA::Symbol) = conjA == :C ? :N : :C
+trivtuple(N) = ntuple(identity, N)
+
+function _repartition(p::IndexTuple, N₁::Int)
+    length(p) >= N₁ ||
+        throw(ArgumentError("cannot repartition $(typeof(p)) to $N₁, $(length(p) - N₁)"))
+    return p[1:N₁], p[(N₁ + 1):end]
+end
+_repartition(p::Index2Tuple, N₁::Int) = _repartition(linearize(p), N₁)
+function _repartition(p::Union{IndexTuple,Index2Tuple}, ::Index2Tuple{N₁}) where {N₁}
+    return _repartition(p, N₁)
+end
+function _repartition(p::Union{IndexTuple,Index2Tuple},
+                      ::AbstractTensorMap{<:Any,N₁}) where {N₁}
+    return _repartition(p, N₁)
+end
+
+TensorKit.block(t::ZeroTangent, c::Sector) = t
+
+# Constructors
+# ------------
+
+@non_differentiable TensorKit.TensorMap(f::Function, storagetype, cod, dom)
+@non_differentiable TensorKit.isomorphism(args...)
+@non_differentiable TensorKit.isometry(args...)
+@non_differentiable TensorKit.unitary(args...)
+
+function ChainRulesCore.rrule(::Type{<:TensorMap}, d::DenseArray, args...)
+    function TensorMap_pullback(Δt)
+        ∂d = convert(Array, Δt)
+        return NoTangent(), ∂d, fill(NoTangent(), length(args))...
+    end
+    return TensorMap(d, args...), TensorMap_pullback
+end
+
+function ChainRulesCore.rrule(::typeof(convert), T::Type{<:Array}, t::AbstractTensorMap)
+    A = convert(T, t)
+    function convert_pullback(ΔA)
+        ∂t = TensorMap(ΔA, codomain(t), domain(t))
+        return NoTangent(), NoTangent(), ∂t
+    end
+    return A, convert_pullback
+end
+
+# Base Linear Algebra
+# -------------------
+
+function ChainRulesCore.rrule(::typeof(+), a::AbstractTensorMap, b::AbstractTensorMap)
+    plus_pullback(Δc) = NoTangent(), Δc, Δc
+    return a + b, plus_pullback
+end
+
+function ChainRulesCore.rrule(::typeof(-), a::AbstractTensorMap, b::AbstractTensorMap)
+    minus_pullback(Δc) = NoTangent(), Δc, -Δc
+    return a - b, minus_pullback
+end
+
+function ChainRulesCore.rrule(::typeof(*), a::AbstractTensorMap, b::AbstractTensorMap)
+    times_pullback(Δc) = NoTangent(), @thunk(Δc * b'), @thunk(a' * Δc)
+    return a * b, times_pullback
+end
+
+function ChainRulesCore.rrule(::typeof(*), a::AbstractTensorMap, b::Number)
+    times_pullback(Δc) = NoTangent(), @thunk(Δc * b'), @thunk(dot(a, Δc))
+    return a * b, times_pullback
+end
+
+function ChainRulesCore.rrule(::typeof(*), a::Number, b::AbstractTensorMap)
+    times_pullback(Δc) = NoTangent(), @thunk(dot(b, Δc)), @thunk(a' * Δc)
+    return a * b, times_pullback
+end
+
+function ChainRulesCore.rrule(::typeof(permute), tsrc::AbstractTensorMap, p::Index2Tuple)
+    function permute_pullback(Δtdst)
+        invp = TensorKit._canonicalize(TupleTools.invperm(linearize(p)), tsrc)
+        return NoTangent(), permute(unthunk(Δtdst), invp), NoTangent()
+    end
+    return permute(tsrc, p), permute_pullback
+end
+
+function ChainRulesCore.rrule(::typeof(scalar), t::AbstractTensorMap)
+    scalar_pullback(Δc) = NoTangent(), fill!(similar(t), unthunk(Δc))
+    return scalar(t), scalar_pullback
+end
+
+# LinearAlgebra
+# -------------
+
+function ChainRulesCore.rrule(::typeof(tr), A::AbstractTensorMap)
+    tr_pullback(Δtr) = NoTangent(), @thunk(Δtr * id(domain(A)))
+    return tr(A), tr_pullback
+end
+
+function ChainRulesCore.rrule(::typeof(adjoint), A::AbstractTensorMap)
+    adjoint_pullback(Δadjoint) = NoTangent(), adjoint(unthunk(Δadjoint))
+    return adjoint(A), adjoint_pullback
+end
+
+function ChainRulesCore.rrule(::typeof(dot), a::AbstractTensorMap, b::AbstractTensorMap)
+    dot_pullback(Δd) = NoTangent(), @thunk(b * Δd'), @thunk(a * Δd)
+    return dot(a, b), dot_pullback
+end
+
+function ChainRulesCore.rrule(::typeof(norm), a::AbstractTensorMap, p)
+    p == 2 || error("currently only implemented for p = 2")
+    n = norm(a, p)
+    norm_pullback(Δn) = NoTangent(), a * (Δn' + Δn) / (n * 2), NoTangent()
+    return n, norm_pullback
+end
+
+# Factorizations
+# --------------
+
+function ChainRulesCore.rrule(::typeof(TensorKit.tsvd!), t::AbstractTensorMap; kwargs...)
+    U, S, V, ϵ = tsvd(t; kwargs...)
+
+    function tsvd!_pullback((ΔU, ΔS, ΔV, Δϵ))
+        ∂t = similar(t)
+        for (c, b) in blocks(∂t)
+            copyto!(b,
+                    svd_rev(block(U, c), block(S, c), block(V, c),
+                            block(ΔU, c), block(ΔS, c), block(ΔV, c)))
+        end
+
+        return NoTangent(), ∂t
+    end
+
+    return (U, S, V, ϵ), tsvd!_pullback
+end
+
+"""
+    svd_rev(U, S, V, ΔU, ΔS, ΔV; tol=eps(real(scalartype(Σ)))^(4 / 5))
+
+Implements the following back propagation formula for the SVD:
+
+```math
+ΔA = UΔSV' + U(J + J')SV' + US(K + K')V' + \\frac{1}{2}US^{-1}(L' - L)V'\\
+J = F ∘ (U'ΔU), \\qquad K = F ∘ (V'ΔV), \\qquad L = I ∘ (V'ΔV)\\
+F_{i ≠ j} = \\frac{1}{s_j^2 - s_i^2}\\
+F_{ii} = 0
+```
+
+# References
+
+Wan, Zhou-Quan, and Shi-Xin Zhang. 2019. “Automatic Differentiation for Complex Valued SVD.” https://doi.org/10.48550/ARXIV.1909.02659.
+"""
+function svd_rev(U::AbstractMatrix, S::AbstractMatrix, V::AbstractMatrix, ΔU, ΔS, ΔV;
+                 atol::Real=0,
+                 rtol::Real=atol > 0 ? 0 : eps(scalartype(S))^(3 / 4))
+    # project out gauge invariance dependence?
+    # ΔU * U + ΔV * V' = 0
+
+    tol = atol > 0 ? atol : rtol * S[1, 1]
+    F = _invert_S²(S, tol)
+    S⁻¹ = pinv(S; atol=tol)
+
+    term = Diagonal(diag(ΔS))
+
+    J = F .* (U' * ΔU)
+    term += (J + J') * S
+    VΔV = (V * ΔV')
+    K = F .* VΔV
+    term += S * (K + K')
+
+    if scalartype(U) <: Complex && !(ΔV isa ZeroTangent) && !(ΔU isa ZeroTangent)
+        L = LinearAlgebra.Diagonal(diag(VΔV))
+        term += 0.5 * S⁻¹ * (L' - L)
+    end
+
+    ΔA = U * term * V
+
+    if size(U, 1) != size(V, 2)
+        UUd = U * U'
+        VdV = V' * V
+        ΔA += (one(UUd) - UUd) * ΔU * S⁻¹ * V + U * S⁻¹ * ΔV * (one(VdV) - VdV)
+    end
+
+    return ΔA
+end
+
+function _invert_S²(S::AbstractMatrix{T}, tol::Real) where {T<:Real}
+    F = similar(S)
+    @inbounds for i in axes(F, 1), j in axes(F, 2)
+        F[i, j] = if i == j
+            zero(T)
+        else
+            sᵢ, sⱼ = S[i, i], S[j, j]
+            1 / (abs(sⱼ - sᵢ) < tol ? tol : sⱼ^2 - sᵢ^2)
+        end
+    end
+    return F
+end
+
+function ChainRulesCore.rrule(::typeof(leftorth!), t::AbstractTensorMap; alg=QRpos())
+    alg isa TensorKit.QR || alg isa TensorKit.QRpos || error("only QR and QRpos supported")
+    Q, R = leftorth(t; alg)
+    leftorth!_pullback((ΔQ, ΔR)) = NoTangent(), qr_pullback!(similar(t), t, Q, R, ΔQ, ΔR)
+    leftorth!_pullback(::Tuple{ZeroTangent,ZeroTangent}) = ZeroTangent()
+    return (Q, R), leftorth!_pullback
+end
+
+function ChainRulesCore.rrule(::typeof(rightorth!), t::AbstractTensorMap; alg=LQpos())
+    alg isa TensorKit.LQ || alg isa TensorKit.LQpos || error("only LQ and LQpos supported")
+    L, Q = rightorth(t; alg)
+    rightorth!_pullback((ΔL, ΔQ)) = NoTangent(), lq_pullback!(similar(t), t, L, Q, ΔL, ΔQ)
+    rightorth!_pullback(::Tuple{ZeroTangent,ZeroTangent}) = ZeroTangent()
+    return (L, Q), rightorth!_pullback
+end
+
+"""
+    copyltu!(A::AbstractMatrix)
+
+Copy the conjugated lower triangular part of `A` to the upper triangular part.
+"""
+function copyltu!(A::AbstractMatrix)
+    m, n = size(A)
+    for i in axes(A, 1)
+        A[i, i] = real(A[i, i])
+        @inbounds for j in (i + 1):n
+            A[i, j] = conj(A[j, i])
+        end
+    end
+    return A
+end
+
+function qr_pullback!(ΔA::AbstractTensorMap{S}, t::AbstractTensorMap{S},
+                      Q::AbstractTensorMap{S}, R::AbstractTensorMap{S}, ΔQ, ΔR) where {S}
+    for (c, b) in blocks(ΔA)
+        qr_pullback!(b, block(t, c), block(Q, c), block(R, c), block(ΔQ, c), block(ΔR, c))
+    end
+    return ΔA
+end
+
+function qr_pullback!(ΔA, A, Q::M, R::M, ΔQ, ΔR) where {M<:AbstractMatrix}
+    m = qr_rank(R)
+    n = size(R, 2)
+
+    if n == m # full rank
+        q = view(Q, :, 1:m)
+        Δq = view(ΔQ, :, 1:m)
+        r = view(R, 1:m, :)
+        Δr = view(ΔR, 1:m, :)
+        ΔA = qr_pullback_fullrank!(ΔA, q, r, Δq, Δr)
+    else
+        q = view(Q, :, 1:m)
+        Δq = view(ΔQ, :, 1:m) + view(A, :, (m + 1):n) * view(ΔR, :, (m + 1):n)'
+        r = view(R, 1:m, 1:m)
+        Δr = view(ΔR, 1:m, 1:m)
+
+        qr_pullback_fullrank!(view(ΔA, :, 1:m), q, r, Δq, Δr)
+        ΔA[:, (m + 1):n] = q * view(ΔR, :, (m + 1):n)
+    end
+
+    return ΔA
+end
+
+function qr_pullback_fullrank!(ΔA, Q, R, ΔQ, ΔR)
+    b = ΔQ + Q * copyltu!(R * ΔR' - ΔQ' * Q)
+    return adjoint!(ΔA, LinearAlgebra.LAPACK.trtrs!('U', 'N', 'N', R, copy(adjoint(b))))
+end
+
+function lq_pullback!(ΔA::AbstractTensorMap{S}, t::AbstractTensorMap{S},
+                      L::AbstractTensorMap{S}, Q::AbstractTensorMap{S}, ΔL, ΔQ) where {S}
+    for (c, b) in blocks(ΔA)
+        lq_pullback!(b, block(t, c), block(L, c), block(Q, c), block(ΔL, c), block(ΔQ, c))
+    end
+    return ΔA
+end
+
+function lq_pullback!(ΔA, A, L::M, Q::M, ΔL, ΔQ) where {M<:AbstractMatrix}
+    m = qr_rank(L)
+    n = size(L, 1)
+
+    if n == m # full rank
+        l = view(L, :, 1:m)
+        Δl = view(ΔL, :, 1:m)
+        q = view(Q, 1:m, :)
+        Δq = view(ΔQ, 1:m, :)
+        ΔA = lq_pullback_fullrank!(ΔA, l, q, Δl, Δq)
+    else
+        l = view(L, 1:m, 1:m)
+        Δl = view(ΔL, 1:m, 1:m)
+        q = view(Q, 1:m, :)
+        Δq = view(ΔQ, 1:m, :) + view(ΔL, (m + 1):n, 1:m)' * view(A, (m + 1):n, :)
+
+        lq_pullback_fullrank!(view(ΔA, 1:m, :), l, q, Δl, Δq)
+        ΔA[(m + 1):n, :] = view(ΔL, (m + 1):n, :) * q
+    end
+
+    return ΔA
+end
+
+function lq_pullback_fullrank!(ΔA, L, Q, ΔL, ΔQ)
+    mul!(ΔA, copyltu!(L' * ΔL - ΔQ * Q'), Q)
+    axpy!(true, ΔQ, ΔA)
+    return LinearAlgebra.LAPACK.trtrs!('L', 'C', 'N', L, ΔA)
+end
+
+function qr_rank(r::AbstractMatrix)
+    Base.require_one_based_indexing(r)
+    m, n = size(r)
+    r₀ = r[1, 1]
+    i = findfirst(x -> abs(x / r₀) < 1e-12, diag(r))
+    return isnothing(i) ? min(m, n) : i - 1
+end
+
+function ChainRulesCore.rrule(::typeof(Base.convert), ::Type{Dict}, t::AbstractTensorMap)
+    out = convert(Dict, t)
+    function convert_pullback(c)
+        if haskey(c, :data) # :data is the only thing for which this dual makes sense
+            dual = copy(out)
+            dual[:data] = c[:data]
+            return (NoTangent(), NoTangent(), convert(TensorMap, dual))
+        else
+            # instead of zero(t) you can also return ZeroTangent(), which is type unstable
+            return (NoTangent(), NoTangent(), zero(t))
+        end
+    end
+    return out, convert_pullback
+end
+function ChainRulesCore.rrule(::typeof(Base.convert), ::Type{TensorMap},
+                              t::Dict{Symbol,Any})
+    return convert(TensorMap, t), v -> (NoTangent(), NoTangent(), convert(Dict, v))
+end
+
+end

src/TensorKit.jl

@@ -117,6 +117,8 @@ using LinearAlgebra: norm, dot, normalize, normalize!, tr,
                      Diagonal, Hermitian
 import Base.Meta
 
+using PackageExtensionCompat
+
 # Auxiliary files
 #-----------------
 include("auxiliary/auxiliary.jl")
@@ -203,4 +205,10 @@ include("planar/planaroperations.jl")
 # deprecations: to be removed in version 1.0 or sooner
 include("auxiliary/deprecate.jl")
 
+# Extensions
+# ----------
+function __init__()
+    @require_extensions
+end
+
 end