Skip to content
New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

More efficient projection in svd pullback #755

Open
wants to merge 6 commits into
base: main
Choose a base branch
from
Open
Changes from all commits
Commits
File filter

Filter by extension

Filter by extension

Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
50 changes: 26 additions & 24 deletions src/rulesets/LinearAlgebra/factorization.jl
Original file line number Diff line number Diff line change
@@ -216,7 +216,7 @@ end
#####

function _svd_pullback(Ȳ::Tangent, F)
∂X = svd_rev(F, Ȳ.U, Ȳ.S, Ȳ.Vt')
∂X = svd_rev(F, Ȳ.U, Ȳ.S, Ȳ.Vt)
return (NoTangent(), ∂X)
end
_svd_pullback(Ȳ::AbstractThunk, F) = _svd_pullback(unthunk(Ȳ), F)
@@ -243,39 +243,41 @@ function rrule(::typeof(getproperty), F::T, x::Symbol) where T <: SVD
return getproperty(F, x), getproperty_svd_pullback
end

# When not `ZeroTangent`s expect `Ū::AbstractMatrix, s̄::AbstractVector, ::AbstractMatrix`
function svd_rev(USV::SVD, Ū, s̄, )
# When not `ZeroTangent`s expect `Ū::AbstractMatrix, s̄::AbstractVector, V̄t::AbstractMatrix`
function svd_rev(USV::SVD, Ū, s̄, V̄t)
# Note: assuming a thin factorization, i.e. svd(A, full=false), which is the default
U = USV.U
s = USV.S
V = USV.V
Vt = USV.Vt

k = length(s)
T = eltype(s)
F = T[i == j ? 1 : inv(@inbounds s[j]^2 - s[i]^2) for i = 1:k, j = 1:k]

# We do a lot of matrix operations here, so we'll try to be memory-friendly and do
# as many of the computations in-place as possible. Benchmarking shows that the in-
# place functions here are significantly faster than their out-of-place, naively
# implemented counterparts, and allocate no additional memory.
Ut = U'
FUᵀŪ = _mulsubtrans!!(Ut*Ū, F) # F .* (UᵀŪ - ŪᵀU)
FVᵀV̄ = _mulsubtrans!!(Vt*V̄, F) # F .* (VᵀV̄ - V̄ᵀV)
ImUUᵀ = _eyesubx!(U*Ut) # I - UUᵀ
ImVVᵀ = _eyesubx!(V*Vt) # I - VVᵀ

S = Diagonal(s)
S̄ = s̄ isa AbstractZero ? s̄ : Diagonal(s̄)

# TODO: consider using MuladdMacro here
Ā = add!!(U * FUᵀŪ * S, ImUUᵀ * (Ū / S)) * Vt
Ā = add!!(Ā, U * S̄ * Vt)
Ā = add!!(Ā, U * add!!(S * FVᵀV̄ * Vt, (S \ V̄') * ImVVᵀ))

UtŪ = U' * Ū
V̄tV = V̄t * Vt'
M = @inbounds T[
if i == j
s̄[i]
else
(s[j] * (UtŪ[i, j] - UtŪ[j, i]) + s[i] * (V̄tV[j, i] - V̄tV[i, j])) /
(s[j]^2 - s[i]^2)
end for i in 1:k, j in 1:k
]

if size(Vt, 1) == size(Vt, 2)
# V is square, VVᵀ = I and therefore V̄ᵀ - V̄ᵀVVᵀ = 0
Ā = (U * M .+ ((Ū .- U * UtŪ) ./ s')) * Vt
else
# If V is not square then U is, so UUᵀ == I and Ū - UUᵀŪ = 0
Ā = U * (M * Vt .+ ((V̄t .- V̄tV * Vt) ./ s))
end
return Ā
end

function svd_rev(USV::SVD, ::AbstractZero, s̄::AbstractVector, ::AbstractZero)
Ā = USV.U * Diagonal(s̄) * USV.Vt
return Ā
end

#####
##### `eigen`
#####
27 changes: 0 additions & 27 deletions src/rulesets/LinearAlgebra/utils.jl
Original file line number Diff line number Diff line change
@@ -1,33 +1,6 @@
# Some utility functions for optimizing linear algebra operations that aren't specific
# to any particular rule definition

# F .* (X - X'), overwrites X if possible
function _mulsubtrans!!(X::AbstractMatrix{<:Real}, F::AbstractMatrix{<:Real})
T = promote_type(eltype(X), eltype(F))
Y = (T <: eltype(X)) ? X : similar(X, T)
k = size(X, 1)
@inbounds for j = 1:k, i = 1:j # Iterate the upper triangle
if i == j
Y[i,i] = zero(T)
else
Y[i,j], Y[j,i] = F[i,j] * (X[i,j] - X[j,i]), F[j,i] * (X[j,i] - X[i,j])
end
end
return Y
end
_mulsubtrans!!(X::AbstractZero, F::AbstractZero) = X
_mulsubtrans!!(X::AbstractZero, F::AbstractMatrix{<:Real}) = X
_mulsubtrans!!(X::AbstractMatrix{<:Real}, F::AbstractZero) = F

# I - X, overwrites X
function _eyesubx!(X::AbstractMatrix)
n, m = size(X)
@inbounds for j = 1:m, i = 1:n
X[i,j] = (i == j) - X[i,j]
end
return X
end

_extract_imag(x) = complex(0, imag(x))

"""
12 changes: 0 additions & 12 deletions test/rulesets/LinearAlgebra/factorization.jl
Original file line number Diff line number Diff line change
@@ -146,18 +146,6 @@ end
@test dX_thunked == dX_unthunked
end
end

@testset "Helper functions" begin
X = randn(10, 10)
Y = randn(10, 10)
@test ChainRules._mulsubtrans!!(copy(X), Y) ≈ Y .* (X - X')
@test ChainRules._eyesubx!(copy(X)) ≈ I - X

Z = randn(Float32, 10, 10)
result = ChainRules._mulsubtrans!!(copy(Z), Y)
@test result ≈ Y .* (Z - Z')
@test eltype(result) == Float64
end
end

@testset "eigendecomposition" begin