Skip to content

Allow not checking hermiticity in eigh functions #321

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.

Sign up for GitHub

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

Open
kshyatt wants to merge 1 commit into
base: main
Choose a base branch
from
Open

Allow not checking hermiticity in eigh functions #321

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
1 commit
Select commit
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
100 changes: 93 additions & 7 deletions ext/MatrixAlgebraKitEnzymeExt/MatrixAlgebraKitEnzymeExt.jl
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -55,7 +55,6 @@ for (f, pb) in (
(qr_compact!, qr_pullback!),
(lq_compact!, lq_pullback!),
(eig_full!, eig_pullback!),
(eigh_full!, eigh_pullback!),
(left_polar!, left_polar_pullback!),
(right_polar!, right_polar_pullback!),
)
Expand Down Expand Up @@ -121,7 +120,6 @@ end
for (f, pf) in (
(:right_polar!, :right_polar_pushforward!),
(:left_polar!, :left_polar_pushforward!),
(:eigh_full!, :eigh_pushforward!),
(:eig_full!, :eig_pushforward!),
(:qr_full!, :qr_pushforward!),
(:lq_full!, :lq_pushforward!),
Expand Down Expand Up @@ -158,6 +156,91 @@ for (f, pf) in (
end
end

function EnzymeRules.augmented_primal(
Comment thread
lkdvos marked this conversation as resolved.
Resolved
config::EnzymeRules.RevConfigWidth{1},
func::Const{typeof(eigh_full!)},
::Type{RT},
A::Annotation,
arg::Annotation{Tuple{TA, TB}},
alg::Const{<:MatrixAlgebraKit.AbstractAlgorithm};
kwargs...
) where {RT, TA, TB}
# A is overwritten in the primal, but not used in the pullback, so we do not need to cache it.
ret = func.val(A.val, arg.val, alg.val; kwargs...)
# if arg.val === ret, the annotation must be Duplicated or DuplicatedNoNeed
# if arg isa Const, ret may still be modified further down the call graph so we should
# copy it to protect ourselves
A_is_arg1 = !isa(A, Const) && has_equal_storage(A.val, arg.val[1])
A_is_arg2 = !isa(A, Const) && has_equal_storage(A.val, arg.val[2])
A_is_arg = A_is_arg1 || A_is_arg2
cache_arg = arg.val !== ret || A_is_arg || EnzymeRules.overwritten(config)[3] ? copy.(ret) : ret
dret = if EnzymeRules.needs_shadow(config) && ((TA == Nothing && TB == Nothing) || isa(arg, Const))
make_zero.(ret)
elseif EnzymeRules.needs_shadow(config)
arg.dval
else
nothing
end
primal = EnzymeRules.needs_primal(config) ? ret : nothing
return EnzymeRules.AugmentedReturn(primal, dret, (cache_arg, dret))
end
function EnzymeRules.reverse(
config::EnzymeRules.RevConfigWidth{1},
func::Const{typeof(eigh_full!)},
::Type{RT},
cache,
A::Annotation,
arg::Annotation{Tuple{TA, TB}},
alg::Const{<:MatrixAlgebraKit.AbstractAlgorithm};
kwargs...
) where {RT, TA, TB}
cache_arg, darg = cache
# A is not used in the pullback; since we have destroyed A, we insert Aval = nothing
# to trigger an error in case the pushfward is modified to directly use A
Aval = nothing
A_is_arg1 = !isa(A, Const) && has_equal_storage(A.dval, arg.dval[1])
A_is_arg2 = !isa(A, Const) && has_equal_storage(A.dval, arg.dval[2])
A_is_arg = A_is_arg1 || A_is_arg2
argval = cache_arg
if !isa(A, Const)
ΔA = A_is_arg ? make_zero(A.dval) : A.dval
eigh_pullback!(ΔA, Aval, argval, darg)
A_is_arg && (A.dval .= ΔA)
end
if !isa(arg, Const)
A_is_arg1 || make_zero!(arg.dval[1])
A_is_arg2 || make_zero!(arg.dval[2])
end
return (nothing, nothing, nothing)
end
function EnzymeRules.forward(
config::EnzymeRules.FwdConfigWidth{1},
func::Const{typeof(eigh_full!)},
::Type{RT},
A::Annotation,
arg::Annotation,
alg::Annotation{<:MatrixAlgebraKit.AbstractAlgorithm};
kwargs...
) where {RT}
A_is_arg1 = !isa(A, Const) && has_equal_storage(A.val, arg.val[1])
A_is_arg2 = !isa(A, Const) && has_equal_storage(A.val, arg.val[2])
A_is_arg = A_is_arg1 || A_is_arg2
eigh_full!(A.val, arg.val, alg.val)
if !isa(A, Const) && !isa(arg, Const)
eigh_pushforward!(A.dval, A.val, arg.val, arg.dval)
end
!A_is_arg && make_zero!(A.dval)
if EnzymeRules.needs_primal(config) && EnzymeRules.needs_shadow(config)
return arg
elseif EnzymeRules.needs_primal(config)
return arg.val
elseif EnzymeRules.needs_shadow(config)
return arg.dval
else
return nothing
end
end

for (f, pb, pf) in (
(qr_null!, qr_null_pullback!, qr_null_pushforward!),
(lq_null!, lq_null_pullback!, lq_null_pushforward!),
Expand Down Expand Up @@ -415,10 +498,11 @@ for (f!, f_full!, pb!, pf!) in (
::Type{RT},
A::Annotation,
D::Annotation{TD},
alg::Annotation{<:MatrixAlgebraKit.AbstractAlgorithm},
alg::Annotation{<:MatrixAlgebraKit.AbstractAlgorithm};
kwargs...
) where {RT, TD}
# A is overwritten in the primal, but not used in the pullback, so we do not need to cache it.
nD, V = $f_full!(A.val, alg.val)
nD, V = $f_full!(A.val, alg.val; kwargs...)
ret = TD == Nothing ? diagview(nD) : copy!(D.val, diagview(nD))
cache_D = (D.val !== ret) || EnzymeRules.overwritten(config)[3] ? copy(ret) : nothing
primal = EnzymeRules.needs_primal(config) ? ret : nothing
Expand All @@ -435,7 +519,8 @@ for (f!, f_full!, pb!, pf!) in (
cache,
A::Annotation{TA},
D::Annotation,
alg::Const{<:MatrixAlgebraKit.AbstractAlgorithm},
alg::Const{<:MatrixAlgebraKit.AbstractAlgorithm};
kwargs...
) where {RT, TA}
cache_D, dD, V = cache
Dval = something(cache_D, D.val)
Expand All @@ -457,10 +542,11 @@ for (f!, f_full!, pb!, pf!) in (
::Type{RT},
A::Annotation{TA},
D::Annotation,
alg::Const{<:MatrixAlgebraKit.AbstractAlgorithm},
alg::Const{<:MatrixAlgebraKit.AbstractAlgorithm};
kwargs...
) where {RT, TA}
A_is_arg = !isa(A, Const) && has_equal_storage(A.dval, D.dval)
DV = $f_full!(A.val, alg.val)
DV = $f_full!(A.val, alg.val; kwargs...)
Dval, V = DV
if !isa(A, Const) && !isa(D, Const)
ΔD = A_is_arg ? make_zero(D.dval) : D.dval
Expand Down
48 changes: 24 additions & 24 deletions src/implementations/eigh.jl
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -17,8 +17,8 @@ function check_hermitian(A; atol::Real = default_hermitian_tol(A), rtol::Real =
return nothing
end

function check_input(::typeof(eigh_full!), A::AbstractMatrix, DV, alg::AbstractAlgorithm)
check_hermitian(A, alg)
function check_input(::typeof(eigh_full!), A::AbstractMatrix, DV, alg::AbstractAlgorithm; check_hermiticity::Bool = true)
check_hermiticity && check_hermitian(A, alg)
D, V = DV
m = size(A, 1)
@assert D isa Diagonal && V isa AbstractMatrix
Expand All @@ -28,17 +28,17 @@ function check_input(::typeof(eigh_full!), A::AbstractMatrix, DV, alg::AbstractA
@check_scalar(V, A)
return nothing
end
function check_input(::typeof(eigh_vals!), A::AbstractMatrix, D, alg::AbstractAlgorithm)
check_hermitian(A, alg)
function check_input(::typeof(eigh_vals!), A::AbstractMatrix, D, alg::AbstractAlgorithm; check_hermiticity::Bool = true)
check_hermiticity && check_hermitian(A, alg)
m = size(A, 1)
@assert D isa AbstractVector
@check_size(D, (m,))
@check_scalar(D, A, real)
return nothing
end

function check_input(::typeof(eigh_full!), A::AbstractMatrix, DV, alg::DiagonalAlgorithm)
check_hermitian(A, alg)
function check_input(::typeof(eigh_full!), A::AbstractMatrix, DV, alg::DiagonalAlgorithm; check_hermiticity::Bool = true)
check_hermiticity && check_hermitian(A, alg)
@assert isdiag(A)
m = size(A, 1)
D, V = DV
Expand All @@ -50,8 +50,8 @@ function check_input(::typeof(eigh_full!), A::AbstractMatrix, DV, alg::DiagonalA
return nothing
end

function check_input(::typeof(eigh_vals!), A::AbstractMatrix, D, alg::DiagonalAlgorithm)
check_hermitian(A, alg)
function check_input(::typeof(eigh_vals!), A::AbstractMatrix, D, alg::DiagonalAlgorithm; check_hermiticity::Bool = true)
check_hermiticity && check_hermitian(A, alg)
@assert isdiag(A)
m = size(A, 1)
@assert D isa AbstractVector
Expand Down Expand Up @@ -87,11 +87,11 @@ end
# DefaultAlgorithm intercepts
# ---------------------------
for f! in (:eigh_full!, :eigh_vals!, :eigh_trunc!, :eigh_trunc_no_error!)
@eval function $f!(A::AbstractMatrix, alg::DefaultAlgorithm)
return $f!(A, select_algorithm($f!, A, nothing; alg.kwargs...))
@eval function $f!(A::AbstractMatrix, alg::DefaultAlgorithm; check_hermiticity::Bool = true)
return $f!(A, select_algorithm($f!, A, nothing; check_hermiticity, alg.kwargs...))
end
@eval function $f!(A::AbstractMatrix, out, alg::DefaultAlgorithm)
return $f!(A, out, select_algorithm($f!, A, nothing; alg.kwargs...))
@eval function $f!(A::AbstractMatrix, out, alg::DefaultAlgorithm; check_hermiticity::Bool = true)
return $f!(A, out, select_algorithm($f!, A, nothing; check_hermiticity, alg.kwargs...))
end
end

Expand Down Expand Up @@ -120,13 +120,13 @@ for (f, f_lapack!, Alg) in (

# MatrixAlgebraKit wrappers
@eval begin
function eigh_full!(A::AbstractMatrix, DV, alg::$Alg)
check_input(eigh_full!, A, DV, alg)
function eigh_full!(A::AbstractMatrix, DV, alg::$Alg; check_hermiticity::Bool = true)
check_input(eigh_full!, A, DV, alg; check_hermiticity)
$eigh_full_f!(A, DV; alg.kwargs...)
return DV
end
function eigh_vals!(A::AbstractMatrix, D, alg::$Alg)
check_input(eigh_vals!, A, D, alg)
function eigh_vals!(A::AbstractMatrix, D, alg::$Alg; check_hermiticity::Bool = true)
check_input(eigh_vals!, A, D, alg; check_hermiticity)
$eigh_vals_f!(A, D; alg.kwargs...)
return D
end
Expand Down Expand Up @@ -162,22 +162,22 @@ for (f, f_lapack!, Alg) in (
end
end

function eigh_trunc!(A, DV, alg::TruncatedAlgorithm)
D, V = eigh_full!(A, DV, alg.alg)
function eigh_trunc!(A, DV, alg::TruncatedAlgorithm; check_hermiticity::Bool = true)
D, V = eigh_full!(A, DV, alg.alg; check_hermiticity)
DVtrunc, ind = truncate(eigh_trunc!, (D, V), alg.trunc)
return DVtrunc..., truncation_error!(diagview(D), ind)
end

function eigh_trunc_no_error!(A, DV, alg::TruncatedAlgorithm)
D, V = eigh_full!(A, DV, alg.alg)
function eigh_trunc_no_error!(A, DV, alg::TruncatedAlgorithm; check_hermiticity::Bool = true)
D, V = eigh_full!(A, DV, alg.alg; check_hermiticity)
DVtrunc, ind = truncate(eigh_trunc!, (D, V), alg.trunc)
return DVtrunc
end

# Diagonal logic
# --------------
function eigh_full!(A::Diagonal, DV, alg::DiagonalAlgorithm)
check_input(eigh_full!, A, DV, alg)
function eigh_full!(A::Diagonal, DV, alg::DiagonalAlgorithm; check_hermiticity::Bool = true)
check_input(eigh_full!, A, DV, alg; check_hermiticity)
D, V = DV
diagA = diagview(A)
I = sortperm(diagA; by = real)
Expand All @@ -193,8 +193,8 @@ function eigh_full!(A::Diagonal, DV, alg::DiagonalAlgorithm)
return D, V
end

function eigh_vals!(A::Diagonal, D, alg::DiagonalAlgorithm)
check_input(eigh_vals!, A, D, alg)
function eigh_vals!(A::Diagonal, D, alg::DiagonalAlgorithm; check_hermiticity::Bool = true)
check_input(eigh_vals!, A, D, alg; check_hermiticity)
Ad = diagview(A)
if D === Ad
sort!(Ad)
Expand Down
Loading