Skip to content

Differentiable matrix factorizations using ImplicitDifferentiation.jl.

License

Notifications You must be signed in to change notification settings

mohamed82008/DifferentiableFactorizations.jl

Folders and files

NameName
Last commit message
Last commit date

Latest commit

 

History

12 Commits
 
 
 
 
 
 
 
 
 
 
 
 
 
 

Repository files navigation

DifferentiableFactorizations

Build Status Coverage

This package contains a bunch of differentiable matrix factorizations differentiated using the implicit function theorem as implemented in ImplicitDifferentiation.jl. The derivatives computed are only correct if the computed factorization of the matrix is unique and differentiable. If the implementation does not guarantee uniqueness and differentiability, the solution reported cannot be trusted. Theoretically in some cases, the factorization may only be unique up to a permutation or a sign flip. In those cases if the implementation guarantees a unique and differentiable output, the derivatives reported are still valid. This can be experimentally tested by comparing against finite difference.

Installation

using Pkg
Pkg.add(url="https://github.com/mohamed82008/DifferentiableFactorizations.jl")
Pkg.add("Zygote")

Loading

using DifferentiableFactorizations, Zygote, LinearAlgebra

Cholesky factorization

A = rand(3, 3)

# Forward pass: A = L * L' or A = U' * U (i.e. L' == U)
(; L, U) = diff_cholesky(A' * A + 2I)

# Differentiation
f(A) = diff_cholesky(A' * A + 2I).L
zjac = Zygote.jacobian(f, A)[1]

LU factorization

A = rand(3, 3)

# Forward pass: A[p, :] = L * U for a permutation vector p
(; L, U, p) = diff_lu(A)

# Differentiation
f(A) = vec(diff_lu(A).U)
zjac = Zygote.jacobian(f, A)[1]

QR factorization

A = rand(3, 2)

# Forward pass: A = Q * R
(; Q, R) = diff_qr(A)

# Differentiation
f(A) = vec(diff_qr(A).Q)
zjac = Zygote.jacobian(f, A)[1]

Singular value decomposition (SVD)

A = rand(3, 2)

# Forward pass: A = U * Diagonal(S) * V'
(; U, S, V) = diff_svd(A)

# Differentiation
f(A) = vec(diff_svd(A).U)
zjac = Zygote.jacobian(f, A)[1]

Eigenvalue decomposition

A = rand(3, 3)

# Forward pass: `s` are the eigenvalues and `V` are the eigenvectors
(; s, V) = diff_eigen(A' * A)

# Differentiation
f(A) = vec(diff_eigen(A' * A).s)
zjac = Zygote.jacobian(f, A)[1]

Generalized eigenvalue decomposition

A = rand(3, 3)
B = rand(3, 3)

# Forward pass: `s` are the eigenvalues and `V` are the eigenvectors
(; s, V) = diff_eigen(A' * A, B' * B + 2I)

# Differentiation
f(B) = vec(diff_eigen(A' * A, B' * B + 2I).V)
zjac = Zygote.jacobian(f, B)[1]

Schur decomposition

A = rand(3, 3)

# Forward pass: Z * T * Z' ≈ A + A'
(; Z, T) = diff_schur(A + A')

# Differentiation
f(A) = vec(diff_schur(A + A').T)
zjac = Zygote.jacobian(f, A)[1]

Generalized Schur decomposition

A = rand(3, 3)
B = rand(3, 3)

# Forward pass: left * S * right' ≈ A + A' and left * T * right' ≈ B + B'
(; left, right, S, T) = diff_schur(A + A', B + B')

# Differentiation
f(B) = vec(diff_schur(A + A', B + B').T)
zjac = Zygote.jacobian(f, B)[1]