det, logdet, logabsdet rrules for SparseMatrixCSC

Author: ElOceanografo

Project.toml

@@ -13,6 +13,7 @@ LinearAlgebra = "37e2e46d-f89d-539d-b4ee-838fcccc9c8e"
 Random = "9a3f8284-a2c9-5f02-9a11-845980a1fd5c"
 RealDot = "c1ae055f-0cd5-4b69-90a6-9a35b1a98df9"
 SparseArrays = "2f01184e-e22b-5df5-ae63-d93ebab69eaf"
+SparseInverseSubset = "dc90abb0-5640-4711-901d-7e5b23a2fada"
 Statistics = "10745b16-79ce-11e8-11f9-7d13ad32a3b2"
 StructArrays = "09ab397b-f2b6-538f-b94a-2f83cf4a842a"
 

src/rulesets/LinearAlgebra/structured.jl

@@ -1,5 +1,6 @@
 # Structured matrices
 using LinearAlgebra: AbstractTriangular
+using SparseInverseSubset
 
 # Matrix wrapper types that we know are square and are thus potentially invertible. For
 # these we can use simpler definitions for `/` and `\`.
@@ -267,3 +268,40 @@ function rrule(::typeof(logdet), X::Union{Diagonal, AbstractTriangular})
     end
     return y, logdet_pullback
 end
+
+function rrule(::typeof(logabsdet), x::SparseMatrixCSC)
+    F = cholesky(x)
+    L, D, U, P = SparseInverseSubset.get_ldup(F)
+    Ω = logabsdet(D)
+    function logabsdet_pullback(ΔΩ)
+        (Δy, Δsigny) = ΔΩ
+        (_, signy) = Ω
+        f = signy' * Δsigny
+        imagf = f - real(f)
+        g = real(Δy) + imagf
+        Z, P = sparseinv(F, depermute=true)
+        ∂x = g * Z'
+        return (NoTangent(), ∂x)
+    end
+    return Ω, logabsdet_pullback
+end
+
+function rrule(::typeof(logdet), x::SparseMatrixCSC)
+    Ω = logdet(x)
+    function logdet_pullback(ΔΩ)
+        Z, p = sparseinv(x, depermute=true)
+        ∂x = x isa Number ? ΔΩ / x' : ΔΩ * Z'
+        return (NoTangent(), ∂x)
+    end
+    return Ω, logdet_pullback
+end
+
+function rrule(::typeof(det), x::SparseMatrixCSC)
+    Ω = det(x)
+    function det_pullback(ΔΩ)
+        Z, _ = sparseinv(x, depermute=true)
+        ∂x = x isa Number ? ΔΩ : Z' * dot(Ω, ΔΩ)
+        return (NoTangent(), ∂x)
+    end
+    return Ω, det_pullback
+end

test/rulesets/LinearAlgebra/structured.jl

@@ -161,4 +161,14 @@
             end
         end
     end
+
+    @testset "[log[abs[det]]] SparseMatrixCSC" begin
+        ii = 1:5
+        jj = 1:5
+        x = ones(5)
+        A = sparse(ii, jj, x)
+        test_rrule(logabsdet, A)
+        test_rrule(logdet, A)
+        test_rrule(det, A)
+    end
 end