search

SciML / LinearSolve.jl

LinearSolve.jl: High-Performance Unified Interface for Linear Solvers in Julia. Easily switch between factorization and Krylov methods, add preconditioners, and all in one interface.
https://docs.sciml.ai/LinearSolve/stable/
Other
241 stars 52 forks source link

Add ForwardDiff rules #434

Open sharanry opened 9 months ago

sharanry commented 9 months ago
codecov[bot] commented 9 months ago

Codecov Report

Attention: 48 lines in your changes are missing coverage. Please review.

Comparison is base (9aaf9b3) 63.96% compared to head (829a914) 27.07%.

Files Patch % Lines
ext/LinearSolveForwardDiff.jl 0.00% 46 Missing :warning:
src/common.jl 0.00% 2 Missing :warning:
Additional details and impacted files ```diff @@ Coverage Diff @@ ## main #434 +/- ## =========================================== - Coverage 63.96% 27.07% -36.89% =========================================== Files 27 28 +1 Lines 2106 2135 +29 =========================================== - Hits 1347 578 -769 - Misses 759 1557 +798 ```

:umbrella: View full report in Codecov by Sentry.
:loudspeaker: Have feedback on the report? Share it here.

sharanry commented 9 months ago

Still taking a look at performance improvements.

Figured out the method dispatch ambiguities for all methods Krylov:

Stack Trace ```julia ERROR: LoadError: MethodError: no method matching solve!(::Krylov.GmresSolver{ForwardDiff.Dual{ForwardDiff.Tag{var"#fb#31"{KrylovJL{typeof(Krylov.gmres!), Int64, Tuple{}, Base.Pairs{Symbol, Union{}, Tuple{}, NamedTuple{(), Tuple{}}}}}, Float64}, Float64, 4}, ForwardDiff.Dual{ForwardDiff.Tag{var"#fb#31"{KrylovJL{typeof(Krylov.gmres!), Int64, Tuple{}, Base.Pairs{Symbol, Union{}, Tuple{}, NamedTuple{(), Tuple{}}}}}, Float64}, Float64, 4}, Vector{ForwardDiff.Dual{ForwardDiff.Tag{var"#fb#31"{KrylovJL{typeof(Krylov.gmres!), Int64, Tuple{}, Base.Pairs{Symbol, Union{}, Tuple{}, NamedTuple{(), Tuple{}}}}}, Float64}, Float64, 4}}}, ::Matrix{Float64}, ::Vector{ForwardDiff.Dual{ForwardDiff.Tag{var"#fb#31"{KrylovJL{typeof(Krylov.gmres!), Int64, Tuple{}, Base.Pairs{Symbol, Union{}, Tuple{}, NamedTuple{(), Tuple{}}}}}, Float64}, Float64, 4}}; M::LinearAlgebra.UniformScaling{Bool}, N::LinearAlgebra.UniformScaling{Bool}, restart::Bool, atol::ForwardDiff.Dual{ForwardDiff.Tag{var"#fb#31"{KrylovJL{typeof(Krylov.gmres!), Int64, Tuple{}, Base.Pairs{Symbol, Union{}, Tuple{}, NamedTuple{(), Tuple{}}}}}, Float64}, Float64, 4}, rtol::ForwardDiff.Dual{ForwardDiff.Tag{var"#fb#31"{KrylovJL{typeof(Krylov.gmres!), Int64, Tuple{}, Base.Pairs{Symbol, Union{}, Tuple{}, NamedTuple{(), Tuple{}}}}}, Float64}, Float64, 4}, itmax::Int64, verbose::Int64, ldiv::Bool, history::Bool) Closest candidates are: solve!(::Krylov.GpmrSolver{T, FC, S}, ::Any, ::Any, ::AbstractVector{FC}, ::AbstractVector{FC}; C, D, E, F, ldiv, gsp, λ, μ, reorthogonalization, atol, rtol, itmax, timemax, verbose, history, callback, iostream) where {T<:AbstractFloat, FC<:Union{Complex{T}, T}, S<:AbstractVector{FC}} got unsupported keyword arguments "M", "N", "restart" @ Krylov ~/.julia/packages/Krylov/jLgPS/src/krylov_solve.jl:46 solve!(::Krylov.GpmrSolver{T, FC, S}, ::Any, ::Any, ::AbstractVector{FC}, ::AbstractVector{FC}, ::AbstractVector, ::AbstractVector; C, D, E, F, ldiv, gsp, λ, μ, reorthogonalization, atol, rtol, itmax, timemax, verbose, history, callback, iostream) where {T<:AbstractFloat, FC<:Union{Complex{T}, T}, S<:AbstractVector{FC}} got unsupported keyword arguments "M", "N", "restart" @ Krylov ~/.julia/packages/Krylov/jLgPS/src/krylov_solve.jl:59 solve!(::Krylov.CrmrSolver{T, FC, S}, ::Any, ::AbstractVector{FC}; N, ldiv, λ, atol, rtol, itmax, timemax, verbose, history, callback, iostream) where {T<:AbstractFloat, FC<:Union{Complex{T}, T}, S<:AbstractVector{FC}} got unsupported keyword arguments "M", "restart" @ Krylov ~/.julia/packages/Krylov/jLgPS/src/krylov_solve.jl:46 ... Stacktrace: [1] solve!(cache::LinearSolve.LinearCache{Matrix{Float64}, Vector{ForwardDiff.Dual{ForwardDiff.Tag{var"#fb#31"{KrylovJL{typeof(Krylov.gmres!), Int64, Tuple{}, Base.Pairs{Symbol, Union{}, Tuple{}, NamedTuple{(), Tuple{}}}}}, Float64}, Float64, 4}}, Vector{ForwardDiff.Dual{ForwardDiff.Tag{var"#fb#31"{KrylovJL{typeof(Krylov.gmres!), Int64, Tuple{}, Base.Pairs{Symbol, Union{}, Tuple{}, NamedTuple{(), Tuple{}}}}}, Float64}, Float64, 4}}, SciMLBase.NullParameters, KrylovJL{typeof(Krylov.gmres!), Int64, Tuple{}, Base.Pairs{Symbol, Union{}, Tuple{}, NamedTuple{(), Tuple{}}}}, Krylov.GmresSolver{ForwardDiff.Dual{ForwardDiff.Tag{var"#fb#31"{KrylovJL{typeof(Krylov.gmres!), Int64, Tuple{}, Base.Pairs{Symbol, Union{}, Tuple{}, NamedTuple{(), Tuple{}}}}}, Float64}, Float64, 4}, ForwardDiff.Dual{ForwardDiff.Tag{var"#fb#31"{KrylovJL{typeof(Krylov.gmres!), Int64, Tuple{}, Base.Pairs{Symbol, Union{}, Tuple{}, NamedTuple{(), Tuple{}}}}}, Float64}, Float64, 4}, Vector{ForwardDiff.Dual{ForwardDiff.Tag{var"#fb#31"{KrylovJL{typeof(Krylov.gmres!), Int64, Tuple{}, Base.Pairs{Symbol, Union{}, Tuple{}, NamedTuple{(), Tuple{}}}}}, Float64}, Float64, 4}}}, IdentityOperator, IdentityOperator, ForwardDiff.Dual{ForwardDiff.Tag{var"#fb#31"{KrylovJL{typeof(Krylov.gmres!), Int64, Tuple{}, Base.Pairs{Symbol, Union{}, Tuple{}, NamedTuple{(), Tuple{}}}}}, Float64}, Float64, 4}, Bool}, alg::KrylovJL{typeof(Krylov.gmres!), Int64, Tuple{}, Base.Pairs{Symbol, Union{}, Tuple{}, NamedTuple{(), Tuple{}}}}; kwargs::Base.Pairs{Symbol, Union{}, Tuple{}, NamedTuple{(), Tuple{}}}) @ LinearSolve ~/code/enzyme_playground/LS_FD/src/iterative_wrappers.jl:256 [2] solve!(cache::LinearSolve.LinearCache{Matrix{Float64}, Vector{ForwardDiff.Dual{ForwardDiff.Tag{var"#fb#31"{KrylovJL{typeof(Krylov.gmres!), Int64, Tuple{}, Base.Pairs{Symbol, Union{}, Tuple{}, NamedTuple{(), Tuple{}}}}}, Float64}, Float64, 4}}, Vector{ForwardDiff.Dual{ForwardDiff.Tag{var"#fb#31"{KrylovJL{typeof(Krylov.gmres!), Int64, Tuple{}, Base.Pairs{Symbol, Union{}, Tuple{}, NamedTuple{(), Tuple{}}}}}, Float64}, Float64, 4}}, SciMLBase.NullParameters, KrylovJL{typeof(Krylov.gmres!), Int64, Tuple{}, Base.Pairs{Symbol, Union{}, Tuple{}, NamedTuple{(), Tuple{}}}}, Krylov.GmresSolver{ForwardDiff.Dual{ForwardDiff.Tag{var"#fb#31"{KrylovJL{typeof(Krylov.gmres!), Int64, Tuple{}, Base.Pairs{Symbol, Union{}, Tuple{}, NamedTuple{(), Tuple{}}}}}, Float64}, Float64, 4}, ForwardDiff.Dual{ForwardDiff.Tag{var"#fb#31"{KrylovJL{typeof(Krylov.gmres!), Int64, Tuple{}, Base.Pairs{Symbol, Union{}, Tuple{}, NamedTuple{(), Tuple{}}}}}, Float64}, Float64, 4}, Vector{ForwardDiff.Dual{ForwardDiff.Tag{var"#fb#31"{KrylovJL{typeof(Krylov.gmres!), Int64, Tuple{}, Base.Pairs{Symbol, Union{}, Tuple{}, NamedTuple{(), Tuple{}}}}}, Float64}, Float64, 4}}}, IdentityOperator, IdentityOperator, ForwardDiff.Dual{ForwardDiff.Tag{var"#fb#31"{KrylovJL{typeof(Krylov.gmres!), Int64, Tuple{}, Base.Pairs{Symbol, Union{}, Tuple{}, NamedTuple{(), Tuple{}}}}}, Float64}, Float64, 4}, Bool}, alg::KrylovJL{typeof(Krylov.gmres!), Int64, Tuple{}, Base.Pairs{Symbol, Union{}, Tuple{}, NamedTuple{(), Tuple{}}}}) @ LinearSolve ~/code/enzyme_playground/LS_FD/src/iterative_wrappers.jl:225 [3] solve!(::LinearSolve.LinearCache{Matrix{Float64}, Vector{ForwardDiff.Dual{ForwardDiff.Tag{var"#fb#31"{KrylovJL{typeof(Krylov.gmres!), Int64, Tuple{}, Base.Pairs{Symbol, Union{}, Tuple{}, NamedTuple{(), Tuple{}}}}}, Float64}, Float64, 4}}, Vector{ForwardDiff.Dual{ForwardDiff.Tag{var"#fb#31"{KrylovJL{typeof(Krylov.gmres!), Int64, Tuple{}, Base.Pairs{Symbol, Union{}, Tuple{}, NamedTuple{(), Tuple{}}}}}, Float64}, Float64, 4}}, SciMLBase.NullParameters, KrylovJL{typeof(Krylov.gmres!), Int64, Tuple{}, Base.Pairs{Symbol, Union{}, Tuple{}, NamedTuple{(), Tuple{}}}}, Krylov.GmresSolver{ForwardDiff.Dual{ForwardDiff.Tag{var"#fb#31"{KrylovJL{typeof(Krylov.gmres!), Int64, Tuple{}, Base.Pairs{Symbol, Union{}, Tuple{}, NamedTuple{(), Tuple{}}}}}, Float64}, Float64, 4}, ForwardDiff.Dual{ForwardDiff.Tag{var"#fb#31"{KrylovJL{typeof(Krylov.gmres!), Int64, Tuple{}, Base.Pairs{Symbol, Union{}, Tuple{}, NamedTuple{(), Tuple{}}}}}, Float64}, Float64, 4}, Vector{ForwardDiff.Dual{ForwardDiff.Tag{var"#fb#31"{KrylovJL{typeof(Krylov.gmres!), Int64, Tuple{}, Base.Pairs{Symbol, Union{}, Tuple{}, NamedTuple{(), Tuple{}}}}}, Float64}, Float64, 4}}}, IdentityOperator, IdentityOperator, ForwardDiff.Dual{ForwardDiff.Tag{var"#fb#31"{KrylovJL{typeof(Krylov.gmres!), Int64, Tuple{}, Base.Pairs{Symbol, Union{}, Tuple{}, NamedTuple{(), Tuple{}}}}}, Float64}, Float64, 4}, Bool}; kwargs::Base.Pairs{Symbol, Union{}, Tuple{}, NamedTuple{(), Tuple{}}}) @ LinearSolve ~/code/enzyme_playground/LS_FD/src/common.jl:218 [4] solve!(::LinearSolve.LinearCache{Matrix{Float64}, Vector{ForwardDiff.Dual{ForwardDiff.Tag{var"#fb#31"{KrylovJL{typeof(Krylov.gmres!), Int64, Tuple{}, Base.Pairs{Symbol, Union{}, Tuple{}, NamedTuple{(), Tuple{}}}}}, Float64}, Float64, 4}}, Vector{ForwardDiff.Dual{ForwardDiff.Tag{var"#fb#31"{KrylovJL{typeof(Krylov.gmres!), Int64, Tuple{}, Base.Pairs{Symbol, Union{}, Tuple{}, NamedTuple{(), Tuple{}}}}}, Float64}, Float64, 4}}, SciMLBase.NullParameters, KrylovJL{typeof(Krylov.gmres!), Int64, Tuple{}, Base.Pairs{Symbol, Union{}, Tuple{}, NamedTuple{(), Tuple{}}}}, Krylov.GmresSolver{ForwardDiff.Dual{ForwardDiff.Tag{var"#fb#31"{KrylovJL{typeof(Krylov.gmres!), Int64, Tuple{}, Base.Pairs{Symbol, Union{}, Tuple{}, NamedTuple{(), Tuple{}}}}}, Float64}, Float64, 4}, ForwardDiff.Dual{ForwardDiff.Tag{var"#fb#31"{KrylovJL{typeof(Krylov.gmres!), Int64, Tuple{}, Base.Pairs{Symbol, Union{}, Tuple{}, NamedTuple{(), Tuple{}}}}}, Float64}, Float64, 4}, Vector{ForwardDiff.Dual{ForwardDiff.Tag{var"#fb#31"{KrylovJL{typeof(Krylov.gmres!), Int64, Tuple{}, Base.Pairs{Symbol, Union{}, Tuple{}, NamedTuple{(), Tuple{}}}}}, Float64}, Float64, 4}}}, IdentityOperator, IdentityOperator, ForwardDiff.Dual{ForwardDiff.Tag{var"#fb#31"{KrylovJL{typeof(Krylov.gmres!), Int64, Tuple{}, Base.Pairs{Symbol, Union{}, Tuple{}, NamedTuple{(), Tuple{}}}}}, Float64}, Float64, 4}, Bool}) @ LinearSolve ~/code/enzyme_playground/LS_FD/src/common.jl:217 [5] solve(::LinearProblem{Nothing, true, Matrix{Float64}, Vector{ForwardDiff.Dual{ForwardDiff.Tag{var"#fb#31"{KrylovJL{typeof(Krylov.gmres!), Int64, Tuple{}, Base.Pairs{Symbol, Union{}, Tuple{}, NamedTuple{(), Tuple{}}}}}, Float64}, Float64, 4}}, SciMLBase.NullParameters, Base.Pairs{Symbol, Union{}, Tuple{}, NamedTuple{(), Tuple{}}}}, ::KrylovJL{typeof(Krylov.gmres!), Int64, Tuple{}, Base.Pairs{Symbol, Union{}, Tuple{}, NamedTuple{(), Tuple{}}}}; kwargs::Base.Pairs{Symbol, Union{}, Tuple{}, NamedTuple{(), Tuple{}}}) @ LinearSolve ~/code/enzyme_playground/LS_FD/src/common.jl:214 [6] solve(::LinearProblem{Nothing, true, Matrix{Float64}, Vector{ForwardDiff.Dual{ForwardDiff.Tag{var"#fb#31"{KrylovJL{typeof(Krylov.gmres!), Int64, Tuple{}, Base.Pairs{Symbol, Union{}, Tuple{}, NamedTuple{(), Tuple{}}}}}, Float64}, Float64, 4}}, SciMLBase.NullParameters, Base.Pairs{Symbol, Union{}, Tuple{}, NamedTuple{(), Tuple{}}}}, ::KrylovJL{typeof(Krylov.gmres!), Int64, Tuple{}, Base.Pairs{Symbol, Union{}, Tuple{}, NamedTuple{(), Tuple{}}}}) @ LinearSolve ~/code/enzyme_playground/LS_FD/src/common.jl:211 [7] (::var"#fb#31"{KrylovJL{typeof(Krylov.gmres!), Int64, Tuple{}, Base.Pairs{Symbol, Union{}, Tuple{}, NamedTuple{(), Tuple{}}}}})(b::Vector{ForwardDiff.Dual{ForwardDiff.Tag{var"#fb#31"{KrylovJL{typeof(Krylov.gmres!), Int64, Tuple{}, Base.Pairs{Symbol, Union{}, Tuple{}, NamedTuple{(), Tuple{}}}}}, Float64}, Float64, 4}}) @ Main ~/code/enzyme_playground/LS_FD/test/forwarddiff.jl:24 [8] vector_mode_dual_eval! @ ~/.julia/packages/ForwardDiff/PcZ48/src/apiutils.jl:24 [inlined] [9] vector_mode_gradient(f::var"#fb#31"{KrylovJL{typeof(Krylov.gmres!), Int64, Tuple{}, Base.Pairs{Symbol, Union{}, Tuple{}, NamedTuple{(), Tuple{}}}}}, x::Vector{Float64}, cfg::ForwardDiff.GradientConfig{ForwardDiff.Tag{var"#fb#31"{KrylovJL{typeof(Krylov.gmres!), Int64, Tuple{}, Base.Pairs{Symbol, Union{}, Tuple{}, NamedTuple{(), Tuple{}}}}}, Float64}, Float64, 4, Vector{ForwardDiff.Dual{ForwardDiff.Tag{var"#fb#31"{KrylovJL{typeof(Krylov.gmres!), Int64, Tuple{}, Base.Pairs{Symbol, Union{}, Tuple{}, NamedTuple{(), Tuple{}}}}}, Float64}, Float64, 4}}}) @ ForwardDiff ~/.julia/packages/ForwardDiff/PcZ48/src/gradient.jl:89 [10] gradient(f::Function, x::Vector{Float64}, cfg::ForwardDiff.GradientConfig{ForwardDiff.Tag{var"#fb#31"{KrylovJL{typeof(Krylov.gmres!), Int64, Tuple{}, Base.Pairs{Symbol, Union{}, Tuple{}, NamedTuple{(), Tuple{}}}}}, Float64}, Float64, 4, Vector{ForwardDiff.Dual{ForwardDiff.Tag{var"#fb#31"{KrylovJL{typeof(Krylov.gmres!), Int64, Tuple{}, Base.Pairs{Symbol, Union{}, Tuple{}, NamedTuple{(), Tuple{}}}}}, Float64}, Float64, 4}}}, ::Val{true}) @ ForwardDiff ~/.julia/packages/ForwardDiff/PcZ48/src/gradient.jl:0 [11] gradient(f::Function, x::Vector{Float64}, cfg::ForwardDiff.GradientConfig{ForwardDiff.Tag{var"#fb#31"{KrylovJL{typeof(Krylov.gmres!), Int64, Tuple{}, Base.Pairs{Symbol, Union{}, Tuple{}, NamedTuple{(), Tuple{}}}}}, Float64}, Float64, 4, Vector{ForwardDiff.Dual{ForwardDiff.Tag{var"#fb#31"{KrylovJL{typeof(Krylov.gmres!), Int64, Tuple{}, Base.Pairs{Symbol, Union{}, Tuple{}, NamedTuple{(), Tuple{}}}}}, Float64}, Float64, 4}}}) @ ForwardDiff ~/.julia/packages/ForwardDiff/PcZ48/src/gradient.jl:17 [12] gradient(f::Function, x::Vector{Float64}) @ ForwardDiff ~/.julia/packages/ForwardDiff/PcZ48/src/gradient.jl:17 [13] top-level scope @ ~/code/enzyme_playground/LS_FD/test/forwarddiff.jl:41 [14] include(fname::String) @ Base.MainInclude ./client.jl:478 [15] top-level scope @ REPL[20]:1 in expression starting at /Users/sharan/code/enzyme_playground/LS_FD/test/forwarddiff.jl:14 ```
ChrisRackauckas commented 8 months ago

See https://github.com/SciML/NonlinearSolve.jl/pull/340. It should be somewhat similar, in that init should build an extended cache.

ChrisRackauckas commented 8 months ago

Note https://github.com/SciML/SciMLBase.jl/issues/558 as a downstream test case.