Add and test rrule for Woodbury

Author: Alex Robson

Project.toml

@@ -1,21 +1,24 @@
 name = "PDMatsExtras"
 uuid = "2c7acb1b-7338-470f-b38f-951d2bcb9193"
 authors = ["Invenia Technical Computing"]
-version = "2.5.1"
+version = "2.5.2"
 
 [deps]
 ChainRulesCore = "d360d2e6-b24c-11e9-a2a3-2a2ae2dbcce4"
+ChainRulesTestUtils = "cdddcdb0-9152-4a09-a978-84456f9df70a"
+FiniteDifferences = "26cc04aa-876d-5657-8c51-4c34ba976000"
 LinearAlgebra = "37e2e46d-f89d-539d-b4ee-838fcccc9c8e"
 PDMats = "90014a1f-27ba-587c-ab20-58faa44d9150"
 SparseArrays = "2f01184e-e22b-5df5-ae63-d93ebab69eaf"
 SuiteSparse = "4607b0f0-06f3-5cda-b6b1-a6196a1729e9"
+Zygote = "e88e6eb3-aa80-5325-afca-941959d7151f"
 
 [compat]
-ChainRulesCore = "0.9.17, 0.10"
+ChainRulesCore = "1"
 Distributions = "0.23, 0.24"
 FiniteDifferences = "0.11, 0.12"
 PDMats = "0.9, 0.10, 0.11"
-Zygote = "0.5.5"
+Zygote = "0.4, 0.5, 0.6"
 julia = "1"
 
 [extras]

src/PDMatsExtras.jl

@@ -12,5 +12,6 @@ export submat
 include("psd_mat.jl")
 include("woodbury_pd_mat.jl")
 include("utils.jl")
+include("chainrules.jl")
 
 end

src/chainrules.jl

@@ -0,0 +1,46 @@
+@non_differentiable validate_woodbury_arguments(A, D, S)
+
+function ChainRulesCore.rrule(
+    ::typeof(*), A::Real, B::WoodburyPDMat{T, TA, TD, TS}
+) where {T, TA, TD, TS}
+    project_A = ProjectTo(A)
+    project_B = ProjectTo(B)
+
+    function times_pullback(ȳ::AbstractMatrix)
+        Ȳ = unthunk(ȳ)
+        Ā = dot(Ȳ, B)
+        B̄ = A' * Ȳ
+        return (
+           NoTangent(),
+           @thunk(project_A(Ā')),
+           @thunk(project_B(B̄)),
+        )
+    end
+
+    function times_pullback(ȳ::Tangent{<:WoodburyPDMat})
+        Ȳ = unthunk(ȳ)
+        Ā = dot(Ȳ.A * Ȳ.D * Ȳ.A' + Ȳ.S, B)
+        B̄ = Ȳ.A * (A' * Ȳ.D) * Ȳ.A' + A' * Ȳ.S
+        return (
+           NoTangent(),
+           @thunk(project_A(Ā')),
+           @thunk(project_B(B̄)),
+        )
+    end
+    return A * B, times_pullback
+end
+
+function ChainRulesCore.ProjectTo(W::WoodburyPDMat)
+    function dW(W̄)
+        Ā(W̄) = ProjectTo(W.A)(collect((W.D * W.A' * W̄' + W.D * W.A' * W̄)'))
+        D̄(W̄) = ProjectTo(W.D)(W.A' * (W̄) * W.A)
+        S̄(W̄) = ProjectTo(W.S)(W̄)
+        return Tangent{typeof(W)}(; A = Ā(W̄), D = D̄(W̄), S = S̄(W̄))
+    end
+    return dW
+end
+
+
+
+
+

src/woodbury_pd_mat.jl

@@ -62,8 +62,6 @@ function validate_woodbury_arguments(A, D, S)
     end
 end
 
-@non_differentiable validate_woodbury_arguments(A, D, S)
-
 function LinearAlgebra.logdet(W::WoodburyPDMat)
     C_S = cholesky(W.S)
     B = C_S.U' \ (W.A * cholesky(W.D).U')
@@ -90,4 +88,6 @@ end
 # NOTE: the parameterisation to scale up the Woodbury matrix is not unique. Here we
 # implement one way to scale it.
 *(a::WoodburyPDMat, c::Real) = WoodburyPDMat(a.A, a.D * c, a.S * c)
-*(c::Real, a::WoodburyPDMat) = a * c
+*(c::Real, a::WoodburyPDMat) = a * c
+*(c::Diagonal{T}, a::WoodburyPDMat) where {T<:Real} = c * Matrix(a)
+*(c1::Diagonal{T}, a::WoodburyPDMat, c2::Diagonal{T}) where {T} = c1 * Matrix(a) * c2

test/chainrules.jl

@@ -0,0 +1,71 @@
+# Create a struct to hold the fields of the Woodbury. 
+# We do this because the Woodbury cannot represent it's own tangent. 
+# I.e. for Y = f(W,...), W̄ = ∂Y / ∂W, is not necessarily a valid Woodbury. 
+# Consider, e.g. the case of a Positive Diagonal matrix
+struct WoodburyLike
+    A
+    D
+    S
+end
+
+# Overwrite the generic to_bec and replace with the almost identical Woodbury specific.
+# This means that in FiniteDifferences, the WoodburyLike matrix is created instead of the Woodbury. 
+# Because the construction is forced there, this would bypass the valdidation checks on the constructor.  
+function FiniteDifferences.to_vec(x::T) where {T<:WoodburyPDMat}
+    val_vecs_and_backs = map(name -> to_vec(getfield(x, name)), fieldnames(T))
+    vals = first.(val_vecs_and_backs)
+    backs = last.(val_vecs_and_backs)
+
+    v, vals_from_vec = to_vec(vals)
+    function structtype_from_vec(v::Vector{<:Real})
+        val_vecs = vals_from_vec(v)
+        values = map((b, v) -> b(v), backs, val_vecs)
+        WoodburyLike(values...)
+    end
+    return v, structtype_from_vec
+end
+
+# Assign some algebra for the WoodburyLike. 
+WoodburyPDMat(S::WoodburyLike) = WoodburyPDMat(S.A, S.D, S.S)
+Base.:*(A::AbstractVecOrMat, B::WoodburyLike) = A * WoodburyPDMat(B)
+Base.:*(A::WoodburyLike, B::AbstractVecOrMat) = WoodburyPDMat(A) * B
+Base.:*(A::Real, B::WoodburyLike) = A * WoodburyPDMat(B)
+Base.:*(A::WoodburyLike, B::Real) = WoodburyPDMat(A) * B
+LinearAlgebra.dot(A, B::WoodburyLike) = dot(A, WoodburyPDMat(B))
+
+@testset "ChainRules" begin
+
+    W = WoodburyPDMat(rand(4,2), Diagonal(rand(2,)), Diagonal(rand(4,)))
+    R = 2.0
+    D = Diagonal(rand(4,))
+
+    @testset "*(Matrix-Woodbury)" begin
+        test_rrule(*, D, W)
+        test_rrule(*, W, D)
+        test_rrule(*, rand(4,4), W)
+    end
+
+    @testset "*(Real-Woodbury" begin
+        @testset "Matrix Tangent" begin
+            ###
+
+            primal = R * W
+            
+            # Matrix Tangent
+            T = rand_tangent(Matrix(primal))
+            f_jvp = j′vp(ChainRulesTestUtils._fdm, x -> Matrix(*(x...)), T, (R, W))[1]
+            R̄ = dot(T, W')
+            W̄ = conj(R) * T
+
+            R̄ ≈ f_jvp[1]
+            (W.D * W.A' * W̄' + W.D * W.A' * W̄) ≈ f_jvp[2].A' # A transpose.
+            Diagonal(W.A' * (W̄) * W.A) ≈ f_jvp[2].D # D
+            Diagonal(W̄) ≈ f_jvp[2].S # S
+
+            # Cannot get this to work. Here the T will be 
+            # T = rand_tangent(primal::WoodburyPDMat) which breaks. 
+            # test_rrule(*, 5.0, W)
+
+        end
+    end
+end

test/runtests.jl

@@ -1,5 +1,7 @@
+using Test
 using PDMatsExtras
 using ChainRulesCore
+using ChainRulesTestUtils
 using Distributions
 using FiniteDifferences
 using LinearAlgebra
@@ -33,6 +35,7 @@ const TEST_MATRICES = Dict(
     include("test_ad.jl")
 
     include("psd_mat.jl")
+    include("chainrules.jl")
     include("woodbury_pd_mat.jl")
     include("utils.jl")
 end