From 16217ed3fcf50cdb24fca110f8374e18f1f632e2 Mon Sep 17 00:00:00 2001 From: Mohamed Tarek Date: Sun, 4 Jun 2023 08:24:56 +1000 Subject: [PATCH 1/8] better direct lin solver and static array support --- Project.toml | 7 ++- examples/0_basic.jl | 6 +- examples/2_nonlinear_solve.jl | 3 +- ext/ImplicitDifferentiationChainRulesExt.jl | 11 ++-- ext/ImplicitDifferentiationForwardDiffExt.jl | 49 ++++++++++++++-- ext/ImplicitDifferentiationStaticArraysExt.jl | 30 ++++++++++ src/ImplicitDifferentiation.jl | 7 ++- src/implicit_function.jl | 56 ++++++++++++++++--- 8 files changed, 145 insertions(+), 24 deletions(-) create mode 100644 ext/ImplicitDifferentiationStaticArraysExt.jl diff --git a/Project.toml b/Project.toml index 2331525..281a837 100644 --- a/Project.toml +++ b/Project.toml @@ -7,6 +7,7 @@ version = "0.5.0-DEV" AbstractDifferentiation = "c29ec348-61ec-40c8-8164-b8c60e9d9f3d" ChainRulesCore = "d360d2e6-b24c-11e9-a2a3-2a2ae2dbcce4" Krylov = "ba0b0d4f-ebba-5204-a429-3ac8c609bfb7" +LinearAlgebra = "37e2e46d-f89d-539d-b4ee-838fcccc9c8e" LinearOperators = "5c8ed15e-5a4c-59e4-a42b-c7e8811fb125" Requires = "ae029012-a4dd-5104-9daa-d747884805df" SimpleUnPack = "ce78b400-467f-4804-87d8-8f486da07d0a" @@ -14,10 +15,12 @@ SimpleUnPack = "ce78b400-467f-4804-87d8-8f486da07d0a" [weakdeps] ChainRulesCore = "d360d2e6-b24c-11e9-a2a3-2a2ae2dbcce4" ForwardDiff = "f6369f11-7733-5829-9624-2563aa707210" +StaticArrays = "90137ffa-7385-5640-81b9-e52037218182" [extensions] ImplicitDifferentiationChainRulesExt = "ChainRulesCore" ImplicitDifferentiationForwardDiffExt = "ForwardDiff" +ImplicitDifferentiationStaticArraysExt = "StaticArrays" [compat] AbstractDifferentiation = "0.5" @@ -40,14 +43,14 @@ ForwardDiff = "f6369f11-7733-5829-9624-2563aa707210" JET = "c3a54625-cd67-489e-a8e7-0a5a0ff4e31b" JuliaFormatter = "98e50ef6-434e-11e9-1051-2b60c6c9e899" Krylov = "ba0b0d4f-ebba-5204-a429-3ac8c609bfb7" -LinearAlgebra = "37e2e46d-f89d-539d-b4ee-838fcccc9c8e" NLsolve = "2774e3e8-f4cf-5e23-947b-6d7e65073b56" Optim = "429524aa-4258-5aef-a3af-852621145aeb" Pkg = "44cfe95a-1eb2-52ea-b672-e2afdf69b78f" Random = "9a3f8284-a2c9-5f02-9a11-845980a1fd5c" SparseArrays = "2f01184e-e22b-5df5-ae63-d93ebab69eaf" +StaticArrays = "90137ffa-7385-5640-81b9-e52037218182" Test = "8dfed614-e22c-5e08-85e1-65c5234f0b40" Zygote = "e88e6eb3-aa80-5325-afca-941959d7151f" [targets] -test = ["Aqua", "ChainRulesCore", "ChainRulesTestUtils", "ComponentArrays", "Documenter", "ForwardDiff", "JET", "JuliaFormatter", "LinearAlgebra", "NLsolve", "Optim", "Pkg", "Random", "SparseArrays", "Test", "Zygote"] +test = ["Aqua", "ChainRulesCore", "ChainRulesTestUtils", "ComponentArrays", "Documenter", "ForwardDiff", "JET", "JuliaFormatter", "NLsolve", "Optim", "Pkg", "Random", "SparseArrays", "StaticArrays", "Test", "Zygote"] diff --git a/examples/0_basic.jl b/examples/0_basic.jl index 3f85b3d..ead5b21 100644 --- a/examples/0_basic.jl +++ b/examples/0_basic.jl @@ -147,11 +147,9 @@ We can even go higher-order by mixing the two packages (forward-over-reverse mod The only technical requirement is to switch the linear solver to something that can handle dual numbers: =# -manual_linear_solver(A, b) = (Matrix(A) \ b, (solved=true,)) +import ImplicitDifferentiation: direct_linear_solver -implicit_higher_order = ImplicitFunction( - forward, conditions; linear_solver=manual_linear_solver -) +implicit_higher_order = ImplicitFunction(forward, conditions, direct_linear_solver) #= Then the Jacobian itself is differentiable. diff --git a/examples/2_nonlinear_solve.jl b/examples/2_nonlinear_solve.jl index 103a910..b08f394 100644 --- a/examples/2_nonlinear_solve.jl +++ b/examples/2_nonlinear_solve.jl @@ -33,7 +33,8 @@ In this case, the optimization problem boils down to the componentwise square ro function forward_nlsolve(x; method) F!(storage, y) = (storage .= y .^ 2 - x) - initial_y = ones(eltype(x), size(x)) + initial_y = similar(x) + initial_y .= 1 result = nlsolve(F!, initial_y; method) return result.zero end diff --git a/ext/ImplicitDifferentiationChainRulesExt.jl b/ext/ImplicitDifferentiationChainRulesExt.jl index dbfbc55..9b2a49c 100644 --- a/ext/ImplicitDifferentiationChainRulesExt.jl +++ b/ext/ImplicitDifferentiationChainRulesExt.jl @@ -5,6 +5,7 @@ using ChainRulesCore: ChainRulesCore, NoTangent, RuleConfig, ZeroTangent, unthun using ImplicitDifferentiation: ImplicitFunction, PullbackMul!, check_solution using LinearOperators: LinearOperator using SimpleUnPack: @unpack +using LinearAlgebra: mul! """ rrule(rc, implicit, x; kwargs...) @@ -27,7 +28,7 @@ function ChainRulesCore.rrule( ::Val{return_byproduct}; kwargs..., ) where {R,return_byproduct} - @unpack conditions, linear_solver = implicit + @unpack conditions, linear_solver, presolver = implicit y, z = implicit(x, Val(true); kwargs...) n, m = length(x), length(y) @@ -40,7 +41,7 @@ function ChainRulesCore.rrule( Aᵀ_op = LinearOperator(R, m, m, false, false, pbmA) Bᵀ_op = LinearOperator(R, n, m, false, false, pbmB) implicit_pullback = ImplicitPullback( - Aᵀ_op, Bᵀ_op, linear_solver, x, Val(return_byproduct) + presolver(Aᵀ_op, x, y), Bᵀ_op, linear_solver, x, Val(return_byproduct) ) if return_byproduct @@ -69,11 +70,11 @@ end function (implicit_pullback_byproduct::ImplicitPullback{true})((dy, _)) @unpack Aᵀ_op, Bᵀ_op, linear_solver, x = implicit_pullback_byproduct R = eltype(x) - - dy_vec = convert(Vector{R}, vec(unthunk(dy))) + dy_vec = convert(AbstractVector{R}, vec(unthunk(dy))) dF_vec, stats = linear_solver(Aᵀ_op, dy_vec) check_solution(linear_solver, stats) - dx_vec = Bᵀ_op * dF_vec + dx_vec = vec(similar(x)) + mul!(dx_vec, Bᵀ_op, dF_vec) dx_vec .*= -1 dx = reshape(dx_vec, size(x)) return (NoTangent(), dx, NoTangent()) diff --git a/ext/ImplicitDifferentiationForwardDiffExt.jl b/ext/ImplicitDifferentiationForwardDiffExt.jl index 99fa2b1..7c4a568 100644 --- a/ext/ImplicitDifferentiationForwardDiffExt.jl +++ b/ext/ImplicitDifferentiationForwardDiffExt.jl @@ -6,10 +6,12 @@ else using ..ForwardDiff: Dual, Partials, jacobian, partials, value end -using AbstractDifferentiation: ForwardDiffBackend, pushforward_function +using AbstractDifferentiation: ForwardDiffBackend +import AbstractDifferentiation: pushforward_function using ImplicitDifferentiation: ImplicitFunction, PushforwardMul!, check_solution using LinearOperators: LinearOperator using SimpleUnPack: @unpack +using LinearAlgebra: mul! """ implicit(x_and_dx::AbstractArray{<:Dual}; kwargs...) @@ -28,7 +30,7 @@ Keyword arguments are given to both `implicit.forward` and `implicit.conditions` function (implicit::ImplicitFunction)( x_and_dx::AbstractArray{Dual{T,R,N}}, ::Val{return_byproduct}=Val(false); kwargs... ) where {T,R,N,return_byproduct} - @unpack conditions, linear_solver = implicit + @unpack conditions, linear_solver, presolver = implicit x = value.(x_and_dx) y, z = implicit(x, Val(true); kwargs...) @@ -39,11 +41,14 @@ function (implicit::ImplicitFunction)( pfB = pushforward_function(backend, _x -> conditions(_x, y, z; kwargs...), x) A_op = LinearOperator(R, m, m, false, false, PushforwardMul!(pfA, size(y))) B_op = LinearOperator(R, m, n, false, false, PushforwardMul!(pfB, size(x))) + A = presolver(A_op, x, y) dy = map(1:N) do k dₖx_vec = vec(partials.(x_and_dx, k)) - dₖy_vec, stats = linear_solver(A_op, B_op * dₖx_vec) - dₖy_vec .*= -1 + Bdx = vec(similar(y)) + mul!(Bdx, B_op, dₖx_vec) + dₖy_vec, stats = linear_solver(A, Bdx) + dₖy_vec = -dₖy_vec check_solution(linear_solver, stats) reshape(dₖy_vec, size(y)) end @@ -62,4 +67,40 @@ function (implicit::ImplicitFunction)( end end +struct ImplicitPushforward{TA,TB,Y,L} + A::TA + B_op::TB + y::Y + linear_solver::L +end + +function pushforward_function( + implicit::ImplicitFunction, x::AbstractArray{R}; kwargs... +) where {R} + @unpack conditions, linear_solver, presolver = implicit + + y, z = implicit(x, Val(true); kwargs...) + n, m = length(x), length(y) + + backend = ForwardDiffBackend() + pfA = pushforward_function(backend, _y -> conditions(x, _y, z; kwargs...), y) + pfB = pushforward_function(backend, _x -> conditions(_x, y, z; kwargs...), x) + A_op = LinearOperator(R, m, m, false, false, PushforwardMul!(pfA, size(y))) + B_op = LinearOperator(R, m, n, false, false, PushforwardMul!(pfB, size(x))) + A = presolver(A_op, x, y) + return ImplicitPushforward(A, B_op, y, linear_solver) +end + +function (pf::ImplicitPushforward)(dx::AbstractVecOrMat) + (; A, B_op, y, linear_solver) = pf + return mapreduce(hcat, eachcol(dx)) do dₖx_vec + Bdx = vec(similar(y)) + mul!(Bdx, B_op, dₖx_vec) + dₖy_vec, stats = linear_solver(A, Bdx) + dₖy_vec = -dₖy_vec + check_solution(linear_solver, stats) + return dₖy_vec + end +end + end diff --git a/ext/ImplicitDifferentiationStaticArraysExt.jl b/ext/ImplicitDifferentiationStaticArraysExt.jl new file mode 100644 index 0000000..48c6ed1 --- /dev/null +++ b/ext/ImplicitDifferentiationStaticArraysExt.jl @@ -0,0 +1,30 @@ +module ImplicitDifferentiationStaticArraysExt + +@static if isdefined(Base, :get_extension) + using StaticArrays: StaticVector, MMatrix, LU +else + using ..StaticArrays: StaticVector, MMatrix, LU +end + +import ImplicitDifferentiation: direct_presolver, default_presolver, auto_linear_solver +using ImplicitDifferentiation: direct_linear_solver +using LinearAlgebra: lu, mul!, I + +auto_linear_solver(A::LU, b) = direct_linear_solver(A, b) + +function direct_presolver(A, ::StaticVector{N1}, y::StaticVector{N2}) where {N1,N2} + _A = zero(MMatrix{N1,N2,eltype(A)}) + for i in 1:size(_A, 2) + v = similar(y) + v .= 0 + v[i] = 1 + mul!(@view(_A[:, i]), A, v) + end + return lu(_A) +end + +function default_presolver(A, x::StaticVector{N1}, y::StaticVector{N2}) where {N1,N2} + return direct_presolver(A, x, y) +end + +end diff --git a/src/ImplicitDifferentiation.jl b/src/ImplicitDifferentiation.jl index 9e2178d..8f45abe 100644 --- a/src/ImplicitDifferentiation.jl +++ b/src/ImplicitDifferentiation.jl @@ -1,8 +1,10 @@ module ImplicitDifferentiation -using AbstractDifferentiation: LazyJacobian, ReverseRuleConfigBackend, lazy_jacobian +using AbstractDifferentiation: + LazyJacobian, ReverseRuleConfigBackend, lazy_jacobian, pushforward_function using Krylov: KrylovStats, gmres using LinearOperators: LinearOperators, LinearOperator +using LinearAlgebra: lu using Requires: @require using SimpleUnPack: @unpack @@ -17,6 +19,9 @@ export ImplicitFunction @require ForwardDiff = "f6369f11-7733-5829-9624-2563aa707210" begin include("../ext/ImplicitDifferentiationForwardDiffExt.jl") end + @require StaticArrays = "90137ffa-7385-5640-81b9-e52037218182" begin + include("../ext/ImplicitDifferentiationStaticArraysExt.jl") + end end end diff --git a/src/implicit_function.jl b/src/implicit_function.jl index a66f504..c77e808 100644 --- a/src/implicit_function.jl +++ b/src/implicit_function.jl @@ -24,6 +24,19 @@ Apply `forward.f` to `x`, returning a dummy byproduct `z(x)=0` if needed. (forward::Forward{true})(x; kwargs...) = forward.f(x; kwargs...) (forward::Forward{false})(x; kwargs...) = (forward.f(x; kwargs...), 0) +auto_linear_solver(A, b) = gmres(A, b) + +direct_linear_solver(A, b) = (A \ b, (; solved=true)) + +function direct_presolver(A, ::T1, ::T2) where {T1<:AbstractArray,T2<:AbstractArray} + return lu(Matrix(A)) +end + +default_presolver(A, ::Any, ::Any) = A + +linear_solver_to_presolver(::typeof(direct_linear_solver)) = direct_presolver +linear_solver_to_presolver(::Any) = default_presolver + """ Conditions{handle_byproduct,C} @@ -59,8 +72,8 @@ Differentiable wrapper for an implicit function defined by a forward mapping and # Constructors - ImplicitFunction(f, c; linear_solver=gmres) - ImplicitFunction(f, c, Val(handle_byproduct); linear_solver=gmres) + ImplicitFunction(f, c; linear_solver=auto_linear_solver) + ImplicitFunction(f, c, Val(handle_byproduct); linear_solver=auto_linear_solver) Construct an `ImplicitFunction` from a forward mapping `f` and conditions `c`, both of which are Julia callables. While `f` does not not need to be compatible with automatic differentiation, `c` has to be. @@ -75,7 +88,7 @@ Given `x ∈ ℝⁿ` and `y ∈ ℝᵈ`, we need as many conditions as output di ∂₂c(x,y(x),z(x)) * ∂y(x) = -∂₁c(x,y(x),z(x)) ``` This requires solving a linear system `A * J = -B`, where `A ∈ ℝᵈˣᵈ`, `B ∈ ℝᵈˣⁿ` and `J ∈ ℝᵈˣⁿ`. -The default linear solver is `Krylov.gmres`, but this can be changed with a keyword argument. +The default linear solver is `Krylov.gmres` when the input is a regular vector and a direct solver when the input is a static vector, but this can be changed with a keyword argument. # Fields @@ -84,21 +97,50 @@ The default linear solver is `Krylov.gmres`, but this can be changed with a keyw - `linear_solver::LS`: a callable of the form `(A,b) -> (u,stats)` such that `Au = b` and `stats.solved ∈ {true,false}`, typically taken from [Krylov.jl](https://github.com/JuliaSmoothOptimizers/Krylov.jl) """ struct ImplicitFunction{ - handle_byproduct,FF<:Forward{handle_byproduct},CC<:Conditions{handle_byproduct},LS + handle_byproduct,FF<:Forward{handle_byproduct},CC<:Conditions{handle_byproduct},LS,P } forward::FF conditions::CC linear_solver::LS + presolver::P function ImplicitFunction( - f, c, ::Val{handle_byproduct}=Val(false); linear_solver=gmres + f, + c, + linear_solver, + presolver=linear_solver_to_presolver(linear_solver), + ::Val{handle_byproduct}=Val(false); ) where {handle_byproduct} forward = Forward{handle_byproduct}(f) conditions = Conditions{handle_byproduct}(c) return new{ - handle_byproduct,typeof(forward),typeof(conditions),typeof(linear_solver) + handle_byproduct, + typeof(forward), + typeof(conditions), + typeof(linear_solver), + typeof(presolver), }( - forward, conditions, linear_solver + forward, conditions, linear_solver, presolver + ) + end + + function ImplicitFunction( + f, c, ::Val{handle_byproduct}=Val(false); + ) where {handle_byproduct} + return ImplicitFunction( + f, c, auto_linear_solver, default_presolver, Val(handle_byproduct) + ) + end + + function ImplicitFunction( + f, c, linear_solver, ::Val{handle_byproduct}; + ) where {handle_byproduct} + return ImplicitFunction( + f, + c, + linear_solver, + linear_solver_to_presolver(linear_solver), + Val(handle_byproduct), ) end end From 7b54c7c19492fd984cc2688ee0c58d0d3cf0d766 Mon Sep 17 00:00:00 2001 From: Guillaume Dalle <22795598+gdalle@users.noreply.github.com> Date: Tue, 20 Jun 2023 12:39:37 +0200 Subject: [PATCH 2/8] Revamped byproduct specification, better tests for static arrays, compatible with static arrays of all sizes --- Project.toml | 1 - examples/0_basic.jl | 4 +- examples/6_byproduct.jl | 2 +- ext/ImplicitDifferentiationChainRulesExt.jl | 61 +++--- ext/ImplicitDifferentiationForwardDiffExt.jl | 75 +++----- ext/ImplicitDifferentiationStaticArraysExt.jl | 38 ++-- src/ImplicitDifferentiation.jl | 7 +- src/conditions.jl | 20 ++ src/forward.jl | 28 +++ src/implicit_function.jl | 175 +++++------------- src/linear_solver.jl | 51 +++++ src/utils.jl | 21 +-- test/misc.jl | 143 -------------- test/runtests.jl | 20 +- test/systematic.jl | 76 ++++++++ 15 files changed, 311 insertions(+), 411 deletions(-) create mode 100644 src/conditions.jl create mode 100644 src/forward.jl create mode 100644 src/linear_solver.jl delete mode 100644 test/misc.jl create mode 100644 test/systematic.jl diff --git a/Project.toml b/Project.toml index 281a837..2dec795 100644 --- a/Project.toml +++ b/Project.toml @@ -24,7 +24,6 @@ ImplicitDifferentiationStaticArraysExt = "StaticArrays" [compat] AbstractDifferentiation = "0.5" -Aqua = "0.6.1" ChainRulesCore = "1.14" ForwardDiff = "0.10" Krylov = "0.8, 0.9" diff --git a/examples/0_basic.jl b/examples/0_basic.jl index ead5b21..20ca4f9 100644 --- a/examples/0_basic.jl +++ b/examples/0_basic.jl @@ -147,9 +147,7 @@ We can even go higher-order by mixing the two packages (forward-over-reverse mod The only technical requirement is to switch the linear solver to something that can handle dual numbers: =# -import ImplicitDifferentiation: direct_linear_solver - -implicit_higher_order = ImplicitFunction(forward, conditions, direct_linear_solver) +implicit_higher_order = ImplicitFunction(forward, conditions, DirectLinearSolver()) #= Then the Jacobian itself is differentiable. diff --git a/examples/6_byproduct.jl b/examples/6_byproduct.jl index b20145b..69f75b7 100644 --- a/examples/6_byproduct.jl +++ b/examples/6_byproduct.jl @@ -1,7 +1,7 @@ # # Byproduct #= -In this example, we show that returning a byproduct with the forward mapping can be useful. +In this example, we show that returning a byproduct with the forward mapping can be useful in certain situations. !!! warning "Work in progress" Come back soon! diff --git a/ext/ImplicitDifferentiationChainRulesExt.jl b/ext/ImplicitDifferentiationChainRulesExt.jl index 9b2a49c..d2af6d1 100644 --- a/ext/ImplicitDifferentiationChainRulesExt.jl +++ b/ext/ImplicitDifferentiationChainRulesExt.jl @@ -1,22 +1,23 @@ module ImplicitDifferentiationChainRulesExt using AbstractDifferentiation: ReverseRuleConfigBackend, pullback_function -using ChainRulesCore: ChainRulesCore, NoTangent, RuleConfig, ZeroTangent, unthunk -using ImplicitDifferentiation: ImplicitFunction, PullbackMul!, check_solution +using ChainRulesCore: ChainRulesCore, NoTangent, RuleConfig, ZeroTangent, rrule, unthunk +using ImplicitDifferentiation: + ImplicitFunction, PullbackMul!, ReturnByproduct, check_solution, presolve, solve +using LinearAlgebra: lmul!, mul! using LinearOperators: LinearOperator using SimpleUnPack: @unpack -using LinearAlgebra: mul! """ rrule(rc, implicit, x; kwargs...) - rrule(rc, implicit, x, Val(return_byproduct); kwargs...) + rrule(rc, implicit, x, ReturnByproduct(); kwargs...) Custom reverse rule for an [`ImplicitFunction`](@ref), to ensure compatibility with reverse mode autodiff. This is only available if ChainRulesCore.jl is loaded (extension). -- If `return_byproduct=false` (the default), this returns a single output `y(x)` with a pullback accepting a single cotangent `̄y`. -- If `return_byproduct=true`, this returns a couple of outputs `(y(x),z(x))` with a pullback accepting a couple of cotangents `(̄y, ̄z)` (remember that `z(x)` is not differentiated so its cotangent is ignored). +- By default, this returns a single output `y(x)` with a pullback accepting a single cotangent `dy`. +- If `ReturnByproduct()` is passed as an argument, this returns a couple of outputs `(y(x),z(x))` with a pullback accepting a couple of cotangents `(dy, dz)` (remember that `z(x)` is not differentiated so its cotangent is ignored). We compute the vector-Jacobian product `Jᵀv` by solving `Aᵀu = v` and setting `Jᵀv = -Bᵀu`. Keyword arguments are given to both `implicit.forward` and `implicit.conditions`. @@ -25,12 +26,12 @@ function ChainRulesCore.rrule( rc::RuleConfig, implicit::ImplicitFunction, x::AbstractArray{R}, - ::Val{return_byproduct}; + ::ReturnByproduct; kwargs..., -) where {R,return_byproduct} - @unpack conditions, linear_solver, presolver = implicit +) where {R} + @unpack conditions, linear_solver = implicit - y, z = implicit(x, Val(true); kwargs...) + y, z = implicit(x, ReturnByproduct(); kwargs...) n, m = length(x), length(y) backend = ReverseRuleConfigBackend(rc) @@ -38,44 +39,40 @@ function ChainRulesCore.rrule( pbB = pullback_function(backend, _x -> conditions(_x, y, z; kwargs...), x) pbmA = PullbackMul!(pbA, size(y)) pbmB = PullbackMul!(pbB, size(y)) + Aᵀ_op = LinearOperator(R, m, m, false, false, pbmA) Bᵀ_op = LinearOperator(R, n, m, false, false, pbmB) - implicit_pullback = ImplicitPullback( - presolver(Aᵀ_op, x, y), Bᵀ_op, linear_solver, x, Val(return_byproduct) - ) + Aᵀ_op_presolved = presolve(linear_solver, Aᵀ_op, y) + + implicit_pullback = ImplicitPullback(Aᵀ_op_presolved, Bᵀ_op, linear_solver, x) - if return_byproduct - return (y, z), implicit_pullback - else - return y, implicit_pullback - end + return (y, z), implicit_pullback end -struct ImplicitPullback{return_byproduct,A,B,L,X} +function ChainRulesCore.rrule( + rc::RuleConfig, implicit::ImplicitFunction, x::AbstractArray{R}; kwargs... +) where {R} + (y, z), implicit_pullback = rrule(rc, implicit, x, ReturnByproduct(); kwargs...) + implicit_pullback_no_byproduct(dy) = implicit_pullback((dy, nothing)) + return y, implicit_pullback_no_byproduct +end + +struct ImplicitPullback{A,B,L,X} Aᵀ_op::A Bᵀ_op::B linear_solver::L x::X - _v::Val{return_byproduct} -end - -function (implicit_pullback_nobyproduct::ImplicitPullback{false})(dy) - @unpack Aᵀ_op, Bᵀ_op, linear_solver, x = implicit_pullback_nobyproduct - implicit_pullback_byproduct = ImplicitPullback( - Aᵀ_op, Bᵀ_op, linear_solver, x, Val(true) - ) - return implicit_pullback_byproduct((dy, nothing)) end -function (implicit_pullback_byproduct::ImplicitPullback{true})((dy, _)) - @unpack Aᵀ_op, Bᵀ_op, linear_solver, x = implicit_pullback_byproduct +function (implicit_pullback::ImplicitPullback)((dy, dz)) + @unpack Aᵀ_op, Bᵀ_op, linear_solver, x = implicit_pullback R = eltype(x) dy_vec = convert(AbstractVector{R}, vec(unthunk(dy))) - dF_vec, stats = linear_solver(Aᵀ_op, dy_vec) + dF_vec, stats = solve(linear_solver, Aᵀ_op, dy_vec) check_solution(linear_solver, stats) dx_vec = vec(similar(x)) mul!(dx_vec, Bᵀ_op, dF_vec) - dx_vec .*= -1 + lmul!(-one(R), dx_vec) dx = reshape(dx_vec, size(x)) return (NoTangent(), dx, NoTangent()) end diff --git a/ext/ImplicitDifferentiationForwardDiffExt.jl b/ext/ImplicitDifferentiationForwardDiffExt.jl index 7c4a568..8c16f39 100644 --- a/ext/ImplicitDifferentiationForwardDiffExt.jl +++ b/ext/ImplicitDifferentiationForwardDiffExt.jl @@ -6,49 +6,51 @@ else using ..ForwardDiff: Dual, Partials, jacobian, partials, value end -using AbstractDifferentiation: ForwardDiffBackend -import AbstractDifferentiation: pushforward_function -using ImplicitDifferentiation: ImplicitFunction, PushforwardMul!, check_solution +using AbstractDifferentiation: + AbstractDifferentiation, ForwardDiffBackend, pushforward_function +using ImplicitDifferentiation: + ImplicitFunction, PushforwardMul!, ReturnByproduct, check_solution, presolve, solve +using LinearAlgebra: lmul!, mul! using LinearOperators: LinearOperator using SimpleUnPack: @unpack -using LinearAlgebra: mul! """ implicit(x_and_dx::AbstractArray{<:Dual}; kwargs...) - implicit(x_and_dx::AbstractArray{<:Dual}, Val(return_byproduct); kwargs...) + implicit(x_and_dx::AbstractArray{<:Dual}, ReturnByproduct(); kwargs...) Overload an [`ImplicitFunction`](@ref) on dual numbers to ensure compatibility with forward mode autodiff. This is only available if ForwardDiff.jl is loaded (extension). -- If `return_byproduct=false` (the default), this returns a single output `y_and_dy(x)`. -- If `return_byproduct=true`, this returns a couple of outputs `(y_and_dy(x),z(x))` (remember that `z(x)` is not differentiated so `dz(x)` doesn't exist). +- By default, this returns a single output `y_and_dy(x)`. +- If `ReturnByproduct()` is passed as an argument, this returns a couple of outputs `(y_and_dy(x),z(x))` (remember that `z(x)` is not differentiated so `dz(x)` doesn't exist). We compute the Jacobian-vector product `Jv` by solving `Au = Bv` and setting `Jv = u`. Keyword arguments are given to both `implicit.forward` and `implicit.conditions`. """ function (implicit::ImplicitFunction)( - x_and_dx::AbstractArray{Dual{T,R,N}}, ::Val{return_byproduct}=Val(false); kwargs... -) where {T,R,N,return_byproduct} - @unpack conditions, linear_solver, presolver = implicit + x_and_dx::AbstractArray{Dual{T,R,N}}, ::ReturnByproduct; kwargs... +) where {T,R,N} + @unpack conditions, linear_solver = implicit x = value.(x_and_dx) - y, z = implicit(x, Val(true); kwargs...) + y, z = implicit(x, ReturnByproduct(); kwargs...) n, m = length(x), length(y) backend = ForwardDiffBackend() pfA = pushforward_function(backend, _y -> conditions(x, _y, z; kwargs...), y) pfB = pushforward_function(backend, _x -> conditions(_x, y, z; kwargs...), x) + A_op = LinearOperator(R, m, m, false, false, PushforwardMul!(pfA, size(y))) B_op = LinearOperator(R, m, n, false, false, PushforwardMul!(pfB, size(x))) - A = presolver(A_op, x, y) + A_op_presolved = presolve(linear_solver, A_op, y) dy = map(1:N) do k dₖx_vec = vec(partials.(x_and_dx, k)) Bdx = vec(similar(y)) mul!(Bdx, B_op, dₖx_vec) - dₖy_vec, stats = linear_solver(A, Bdx) - dₖy_vec = -dₖy_vec + dₖy_vec, stats = solve(linear_solver, A_op_presolved, Bdx) + lmul!(-one(R), dₖy_vec) check_solution(linear_solver, stats) reshape(dₖy_vec, size(y)) end @@ -60,47 +62,14 @@ function (implicit::ImplicitFunction)( reshape(y_and_dy_vec, size(y)) end - if return_byproduct - return y_and_dy, z - else - return y_and_dy - end + return y_and_dy, z end -struct ImplicitPushforward{TA,TB,Y,L} - A::TA - B_op::TB - y::Y - linear_solver::L -end - -function pushforward_function( - implicit::ImplicitFunction, x::AbstractArray{R}; kwargs... -) where {R} - @unpack conditions, linear_solver, presolver = implicit - - y, z = implicit(x, Val(true); kwargs...) - n, m = length(x), length(y) - - backend = ForwardDiffBackend() - pfA = pushforward_function(backend, _y -> conditions(x, _y, z; kwargs...), y) - pfB = pushforward_function(backend, _x -> conditions(_x, y, z; kwargs...), x) - A_op = LinearOperator(R, m, m, false, false, PushforwardMul!(pfA, size(y))) - B_op = LinearOperator(R, m, n, false, false, PushforwardMul!(pfB, size(x))) - A = presolver(A_op, x, y) - return ImplicitPushforward(A, B_op, y, linear_solver) -end - -function (pf::ImplicitPushforward)(dx::AbstractVecOrMat) - (; A, B_op, y, linear_solver) = pf - return mapreduce(hcat, eachcol(dx)) do dₖx_vec - Bdx = vec(similar(y)) - mul!(Bdx, B_op, dₖx_vec) - dₖy_vec, stats = linear_solver(A, Bdx) - dₖy_vec = -dₖy_vec - check_solution(linear_solver, stats) - return dₖy_vec - end +function (implicit::ImplicitFunction)( + x_and_dx::AbstractArray{Dual{T,R,N}}; kwargs... +) where {T,R,N} + y_and_dy, z = implicit(x_and_dx, ReturnByproduct(); kwargs...) + return y_and_dy end end diff --git a/ext/ImplicitDifferentiationStaticArraysExt.jl b/ext/ImplicitDifferentiationStaticArraysExt.jl index 48c6ed1..629be7c 100644 --- a/ext/ImplicitDifferentiationStaticArraysExt.jl +++ b/ext/ImplicitDifferentiationStaticArraysExt.jl @@ -1,30 +1,32 @@ module ImplicitDifferentiationStaticArraysExt @static if isdefined(Base, :get_extension) - using StaticArrays: StaticVector, MMatrix, LU + using StaticArrays: StaticArray, MMatrix else - using ..StaticArrays: StaticVector, MMatrix, LU + using ..StaticArrays: StaticArray, MMatrix end -import ImplicitDifferentiation: direct_presolver, default_presolver, auto_linear_solver -using ImplicitDifferentiation: direct_linear_solver -using LinearAlgebra: lu, mul!, I +import ImplicitDifferentiation: ImplicitDifferentiation, DirectLinearSolver +using LinearAlgebra: lu, mul! -auto_linear_solver(A::LU, b) = direct_linear_solver(A, b) +_prodsize(::Type{Tuple{}}) = 1 +_prodsize(::Type{Tuple{N1}}) where {N1} = N1 +_prodsize(::Type{Tuple{N1,N2}}) where {N1,N2} = N1 * N2 +_prodsize(::Type{Tuple{N1,N2,N3}}) where {N1,N2,N3} = N1 * N2 * N3 +_prodsize(::Type{Tuple{N1,N2,N3,N4}}) where {N1,N2,N3,N4} = N1 * N2 * N3 * N4 -function direct_presolver(A, ::StaticVector{N1}, y::StaticVector{N2}) where {N1,N2} - _A = zero(MMatrix{N1,N2,eltype(A)}) - for i in 1:size(_A, 2) - v = similar(y) - v .= 0 - v[i] = 1 - mul!(@view(_A[:, i]), A, v) +function ImplicitDifferentiation.presolve( + ::DirectLinearSolver, A, y::StaticArray{S,T,N} +) where {S,T,N} + S_prod = _prodsize(S) + A_static = zero(MMatrix{S_prod,S_prod,T}) + for i in axes(A_static, 2) + v = vec(similar(y)) + v .= zero(T) + v[i] = one(T) + mul!(@view(A_static[:, i]), A, v) end - return lu(_A) -end - -function default_presolver(A, x::StaticVector{N1}, y::StaticVector{N2}) where {N1,N2} - return direct_presolver(A, x, y) + return lu(A_static) end end diff --git a/src/ImplicitDifferentiation.jl b/src/ImplicitDifferentiation.jl index 8f45abe..dab8038 100644 --- a/src/ImplicitDifferentiation.jl +++ b/src/ImplicitDifferentiation.jl @@ -1,7 +1,5 @@ module ImplicitDifferentiation -using AbstractDifferentiation: - LazyJacobian, ReverseRuleConfigBackend, lazy_jacobian, pushforward_function using Krylov: KrylovStats, gmres using LinearOperators: LinearOperators, LinearOperator using LinearAlgebra: lu @@ -9,9 +7,14 @@ using Requires: @require using SimpleUnPack: @unpack include("utils.jl") +include("forward.jl") +include("conditions.jl") +include("linear_solver.jl") include("implicit_function.jl") export ImplicitFunction +export IterativeLinearSolver, DirectLinearSolver +export HandleByproduct, ReturnByproduct @static if !isdefined(Base, :get_extension) include("../ext/ImplicitDifferentiationChainRulesExt.jl") diff --git a/src/conditions.jl b/src/conditions.jl new file mode 100644 index 0000000..d0960bc --- /dev/null +++ b/src/conditions.jl @@ -0,0 +1,20 @@ +""" + Conditions{byproduct,C} + +Callable wrapper for conditions `c::C`, which ensures that a byproduct `z` is always accepted in addition to `x` and `y`. + +The type parameter `byproduct` is a boolean stating whether or not `c` natively accepts `z`. +""" +struct Conditions{byproduct,C} + c::C + function Conditions{byproduct}(c::C) where {byproduct,C} + return new{byproduct,C}(c) + end +end + +function Base.show(io::IO, conditions::Conditions{byproduct}) where {byproduct} + return print(io, "Conditions{$byproduct}($(conditions.c))") +end + +(conditions::Conditions{true})(x, y, z; kwargs...) = conditions.c(x, y, z; kwargs...) +(conditions::Conditions{false})(x, y, z; kwargs...) = conditions.c(x, y; kwargs...) diff --git a/src/forward.jl b/src/forward.jl new file mode 100644 index 0000000..06bb62c --- /dev/null +++ b/src/forward.jl @@ -0,0 +1,28 @@ +""" + Forward{byproduct,F} + +Callable wrapper for a forward mapping `f::F`, which ensures that a byproduct `z(x)` is always returned in addition to `y(x)`. + +The type parameter `byproduct` is a boolean stating whether or not `f` natively returns `z(x)`. +""" +struct Forward{byproduct,F} + f::F + function Forward{byproduct}(f::F) where {byproduct,F} + return new{byproduct,F}(f) + end +end + +function Base.show(io::IO, forward::Forward{byproduct}) where {byproduct} + return print(io, "Forward{$byproduct}($(forward.f))") +end + +function (forward::Forward{true})(x; kwargs...) + y, z = forward.f(x; kwargs...) + return y, z +end + +function (forward::Forward{false})(x; kwargs...) + y = forward.f(x; kwargs...) + z = 0 + return y, z +end diff --git a/src/implicit_function.jl b/src/implicit_function.jl index c77e808..4840007 100644 --- a/src/implicit_function.jl +++ b/src/implicit_function.jl @@ -1,177 +1,88 @@ """ - Forward{handle_byproduct,F} - -Callable wrapper for a forward mapping `f::F`, which ensures that a byproduct `z(x)` is always returned in addition to `y(x)`. - -The type parameter `handle_byproduct` is a boolean stating whether or not `f` natively returns `z(x)`. -""" -struct Forward{handle_byproduct,F} - f::F - function Forward{handle_byproduct}(f::F) where {handle_byproduct,F} - return new{handle_byproduct,F}(f) - end -end - -function Base.show(io::IO, forward::Forward{handle_byproduct}) where {handle_byproduct} - return print("Forward{$handle_byproduct}($(forward.f))") -end - -""" - forward(x; kwargs...) - -Apply `forward.f` to `x`, returning a dummy byproduct `z(x)=0` if needed. -""" -(forward::Forward{true})(x; kwargs...) = forward.f(x; kwargs...) -(forward::Forward{false})(x; kwargs...) = (forward.f(x; kwargs...), 0) - -auto_linear_solver(A, b) = gmres(A, b) - -direct_linear_solver(A, b) = (A \ b, (; solved=true)) - -function direct_presolver(A, ::T1, ::T2) where {T1<:AbstractArray,T2<:AbstractArray} - return lu(Matrix(A)) -end - -default_presolver(A, ::Any, ::Any) = A - -linear_solver_to_presolver(::typeof(direct_linear_solver)) = direct_presolver -linear_solver_to_presolver(::Any) = default_presolver - -""" - Conditions{handle_byproduct,C} - -Callable wrapper for conditions `c::C`, which ensures that a byproduct `z` is always accepted in addition to `x` and `y`. - -The type parameter `handle_byproduct` is a boolean stating whether or not `c` natively accepts `z`. -""" -struct Conditions{handle_byproduct,C} - c::C - function Conditions{handle_byproduct}(c::C) where {handle_byproduct,C} - return new{handle_byproduct,C}(c) - end -end - -function Base.show( - io::IO, conditions::Conditions{handle_byproduct} -) where {handle_byproduct} - return print("Conditions{$handle_byproduct}($(conditions.c))") -end - -""" - conditions(x, y, z; kwargs...) - -Apply `conditions.c` to `(x, y, z)`, discarding `z` beforehand if needed. -""" -(conditions::Conditions{true})(x, y, z; kwargs...) = conditions.c(x, y, z; kwargs...) -(conditions::Conditions{false})(x, y, z; kwargs...) = conditions.c(x, y; kwargs...) - -""" - ImplicitFunction{handle_byproduct,FF<:Forward,CC<:Conditions,LS} + ImplicitFunction{FF,CC,LS} Differentiable wrapper for an implicit function defined by a forward mapping and a set of conditions. # Constructors - ImplicitFunction(f, c; linear_solver=auto_linear_solver) - ImplicitFunction(f, c, Val(handle_byproduct); linear_solver=auto_linear_solver) + ImplicitFunction(f, c) + ImplicitFunction(f, c, linear_solver) + ImplicitFunction(f, c, HandleByproduct()) + ImplicitFunction(f, c, linear_solver, HandleByproduct()) Construct an `ImplicitFunction` from a forward mapping `f` and conditions `c`, both of which are Julia callables. While `f` does not not need to be compatible with automatic differentiation, `c` has to be. # Details -- If `handle_byproduct=false` (the default), the forward mapping is `x -> y(x)` and the conditions are `c(x,y(x)) = 0`. -- If `handle_byproduct=true`, the forward mapping is `x -> (y(x),z(x))` and the conditions are `c(x,y(x),z(x)) = 0`. In this case, `z(x)` can contain additional information generated by the forward mapping, but beware that we consider it constant for differentiation purposes. +- By default, we assume that the forward mapping is `x -> y(x)` and the conditions are `c(x,y(x)) = 0`. +- If `handleByproduct()` is passed as an argument to the constructor, we assume instead that the forward mapping is `x -> (y(x),z(x))` and the conditions are `c(x,y(x),z(x)) = 0`. In this case, `z(x)` can contain additional information generated by the forward mapping, but beware that we consider it constant for differentiation purposes. Given `x ∈ ℝⁿ` and `y ∈ ℝᵈ`, we need as many conditions as output dimensions: `c(x,y,z) ∈ ℝᵈ`. We can then compute the Jacobian of `y(⋅)` using the implicit function theorem: ``` ∂₂c(x,y(x),z(x)) * ∂y(x) = -∂₁c(x,y(x),z(x)) ``` This requires solving a linear system `A * J = -B`, where `A ∈ ℝᵈˣᵈ`, `B ∈ ℝᵈˣⁿ` and `J ∈ ℝᵈˣⁿ`. -The default linear solver is `Krylov.gmres` when the input is a regular vector and a direct solver when the input is a static vector, but this can be changed with a keyword argument. # Fields -- `forward::FF`: a wrapper of type [`Forward`](@ref) coherent with the value of `handle_byproduct` -- `conditions::FF`: a wrapper of type [`Conditions`](@ref) coherent with the value of `handle_byproduct` -- `linear_solver::LS`: a callable of the form `(A,b) -> (u,stats)` such that `Au = b` and `stats.solved ∈ {true,false}`, typically taken from [Krylov.jl](https://github.com/JuliaSmoothOptimizers/Krylov.jl) +- `forward::FF`: a wrapper of type [`Forward`](@ref) around the callable `f` +- `conditions::CC`: a wrapper of type [`Conditions`](@ref) around the callable `c` +- `linear_solver::LS`: an object subtyping [`AbstractLinearSolver`](@ref) (defaults to [`IterativeLinearSolver`](@ref)). """ -struct ImplicitFunction{ - handle_byproduct,FF<:Forward{handle_byproduct},CC<:Conditions{handle_byproduct},LS,P -} +struct ImplicitFunction{FF<:Forward,CC<:Conditions,LS<:AbstractLinearSolver} forward::FF conditions::CC linear_solver::LS - presolver::P - - function ImplicitFunction( - f, - c, - linear_solver, - presolver=linear_solver_to_presolver(linear_solver), - ::Val{handle_byproduct}=Val(false); - ) where {handle_byproduct} - forward = Forward{handle_byproduct}(f) - conditions = Conditions{handle_byproduct}(c) - return new{ - handle_byproduct, - typeof(forward), - typeof(conditions), - typeof(linear_solver), - typeof(presolver), - }( - forward, conditions, linear_solver, presolver - ) - end - function ImplicitFunction( - f, c, ::Val{handle_byproduct}=Val(false); - ) where {handle_byproduct} - return ImplicitFunction( - f, c, auto_linear_solver, default_presolver, Val(handle_byproduct) + function ImplicitFunction(f, c, linear_solver::AbstractLinearSolver) + forward = Forward{false}(f) + conditions = Conditions{false}(c) + return new{typeof(forward),typeof(conditions),typeof(linear_solver)}( + forward, conditions, linear_solver ) end - function ImplicitFunction( - f, c, linear_solver, ::Val{handle_byproduct}; - ) where {handle_byproduct} - return ImplicitFunction( - f, - c, - linear_solver, - linear_solver_to_presolver(linear_solver), - Val(handle_byproduct), + function ImplicitFunction(f, c, linear_solver::AbstractLinearSolver, ::HandleByproduct) + forward = Forward{true}(f) + conditions = Conditions{true}(c) + return new{typeof(forward),typeof(conditions),typeof(linear_solver)}( + forward, conditions, linear_solver ) end end -function Base.show( - io::IO, implicit::ImplicitFunction{handle_byproduct} -) where {handle_byproduct} +function ImplicitFunction(f, c) + linear_solver = IterativeLinearSolver() + return ImplicitFunction(f, c, linear_solver) +end + +function ImplicitFunction(f, c, ::HandleByproduct) + linear_solver = IterativeLinearSolver() + return ImplicitFunction(f, c, linear_solver, HandleByproduct()) +end + +function Base.show(io::IO, implicit::ImplicitFunction) @unpack forward, conditions, linear_solver = implicit - @unpack f = forward - @unpack c = conditions - return print(io, "ImplicitFunction{$handle_byproduct}($f, $c, $linear_solver)") + return print(io, "ImplicitFunction($(forward.f), $(conditions.c), $linear_solver)") end """ implicit(x::AbstractArray; kwargs...) - implicit(x::AbstractArray, Val(return_byproduct), ; kwargs...) Make an [`ImplicitFunction`](@ref) callable by applying the forward mapping `implicit.forward`. +""" +function (implicit::ImplicitFunction)(x::AbstractArray, ; kwargs...) + y, z = implicit.forward(x, ; kwargs...) + return y +end -- If `return_byproduct=false` (the default), this returns a single output `y(x)`. -- If `return_byproduct=true`, this returns a couple of outputs `(y(x),z(x))`. +""" + implicit(x::AbstractArray, ReturnByproduct(); kwargs...) -The argument `return_byproduct` is independent from the type parameter `handle_byproduct` in `ImplicitFunction`, so any combination is possible. +Make an [`ImplicitFunction`](@ref) callable by applying the forward mapping `implicit.forward` and returning the byproduct. """ -function (implicit::ImplicitFunction)( - x::AbstractArray, ::Val{return_byproduct}=Val(false); kwargs... -) where {return_byproduct} +function (implicit::ImplicitFunction)(x::AbstractArray, ::ReturnByproduct; kwargs...) y, z = implicit.forward(x, ; kwargs...) - if return_byproduct - return (y, z) - else - return y - end + return (y, z) end diff --git a/src/linear_solver.jl b/src/linear_solver.jl new file mode 100644 index 0000000..06a813f --- /dev/null +++ b/src/linear_solver.jl @@ -0,0 +1,51 @@ +""" + AbstractLinearSolver + +All linear solvers used within an `ImplicitFunction` must satisfy this interface. + +# Required methods + +- `presolve(linear_solver, A, y)`: return a matrix-like object `A` for which it is cheaper to solve several linear systems with different vectors `b` (a typical example would be to perform LU factorization). +- `solve(linear_solver, A, b)`: return a tuple `(x, stats)` where `x` satisfies `Ax = b` and `stats.solved ∈ {true, false}`. +""" +abstract type AbstractLinearSolver end + +""" + IterativeLinearSolver + +An implementation of `AbstractLinearSolver` using `Krylov.gmres`. +""" +struct IterativeLinearSolver <: AbstractLinearSolver end + +presolve(::IterativeLinearSolver, A, y) = A +solve(::IterativeLinearSolver, A, b) = gmres(A, b) + +""" + DirectLinearSolver + +An implementation of `AbstractLinearSolver` using the built-in `\` operator. +""" +struct DirectLinearSolver <: AbstractLinearSolver end + +presolve(::DirectLinearSolver, A, y) = lu(Matrix(A)) +solve(::DirectLinearSolver, A, b) = (A \ b, (; solved=true)) + +struct SolverFailureException{A,B} <: Exception + solver::A + stats::B +end + +function Base.show(io::IO, sfe::SolverFailureException) + return println( + io, + "SolverFailureException: \n Linear solver: $(sfe.solver) \n Solver stats: $(string(sfe.stats))", + ) +end + +function check_solution(solver, stats) + if stats.solved + return nothing + else + throw(SolverFailureException(solver, stats)) + end +end diff --git a/src/utils.jl b/src/utils.jl index 2f005c6..6ceb647 100644 --- a/src/utils.jl +++ b/src/utils.jl @@ -1,22 +1,5 @@ -struct SolverFailureException{A,B} <: Exception - solver::A - stats::B -end - -function Base.show(io::IO, sfe::SolverFailureException) - return println( - io, - "SolverFailureException: \n Solver: $(sfe.solver) \n Solver stats: $(string(sfe.stats))", - ) -end - -function check_solution(solver, stats) - if stats.solved - return nothing - else - throw(SolverFailureException(solver, stats)) - end -end +struct HandleByproduct end +struct ReturnByproduct end """ PushforwardMul!{P,N} diff --git a/test/misc.jl b/test/misc.jl deleted file mode 100644 index 4fb708d..0000000 --- a/test/misc.jl +++ /dev/null @@ -1,143 +0,0 @@ -using ChainRulesCore -using ChainRulesTestUtils -using ForwardDiff -using ImplicitDifferentiation -using JET -using LinearAlgebra -using Random -using Test -using Zygote - -Random.seed!(63); - -""" - mysqrt(x) - -Compute the elementwise square root, breaking Zygote.jl and ForwardDiff.jl in the process. -""" -function mysqrt(x::AbstractArray) - a = [0.0] - a[1] = first(x) - return sqrt.(x) -end - -myval(::Val{X}) where {X} = X - -function make_implicit_sqrt(::Val{handle_byproduct}) where {handle_byproduct} - if handle_byproduct - forward_byproduct(x) = (mysqrt(x), 0) - conditions_byproduct(x, y, z) = y .^ 2 .- x - implicit = ImplicitFunction(forward_byproduct, conditions_byproduct, Val(true)) - else - forward(x) = mysqrt(x) - conditions(x, y) = y .^ 2 .- x - implicit = ImplicitFunction(forward, conditions) - end - return implicit -end - -for handle_byproduct in (Val(true), Val(false)) - testsetname = myval(handle_byproduct) ? "With byproduct" : "Without byproduct" - @testset "$testsetname" verbose = true begin - implicit = make_implicit_sqrt(handle_byproduct) - # Skipped because of https://github.com/JuliaDiff/ChainRulesTestUtils.jl/issues/232 and because it detects weird type instabilities - @testset verbose = true "ChainRulesTestUtils.jl" begin - @test_skip test_rrule(implicit, x) - @test_skip test_rrule(implicit, X) - end - - @testset verbose = true "Vectors" begin - x = rand(2) - y = implicit(x) - J = Diagonal(0.5 ./ sqrt.(x)) - - @testset "Call" begin - @test (@inferred implicit(x)) ≈ sqrt.(x) - if VERSION >= v"1.7" - test_opt(implicit, (typeof(x),)) - end - end - - @testset verbose = true "Forward" begin - @test ForwardDiff.jacobian(implicit, x) ≈ J - x_and_dx = ForwardDiff.Dual.(x, ((0, 0),)) - for return_byproduct in (true, false) - res_and_dres = @inferred implicit(x_and_dx, Val(return_byproduct)) - if return_byproduct - y_and_dy, z = res_and_dres - @test size(y_and_dy) == size(y) - else - y_and_dy = res_and_dres - @test size(y_and_dy) == size(y) - end - end - end - - @testset "Reverse" begin - @test Zygote.jacobian(implicit, x)[1] ≈ J - for return_byproduct in (true, false) - _, pullback = @inferred rrule( - Zygote.ZygoteRuleConfig(), implicit, x, Val(return_byproduct) - ) - dy, dz = zero(implicit(x)), 0 - if return_byproduct - @test (@inferred pullback((dy, dz))) == pullback((dy, dz)) - _, dx = pullback((dy, dz)) - @test size(dx) == size(x) - else - @test (@inferred pullback(dy)) == pullback(dy) - _, dx = pullback(dy) - @test size(dx) == size(x) - end - end - end - end - - @testset verbose = true "Arrays" begin - X = rand(2, 3, 4) - Y = implicit(X) - JJ = Diagonal(0.5 ./ sqrt.(vec(X))) - - @testset "Call" begin - @test (@inferred implicit(X)) ≈ sqrt.(X) - if VERSION >= v"1.7" - test_opt(implicit, (typeof(X),)) - end - end - - @testset "Forward" begin - @test ForwardDiff.jacobian(implicit, X) ≈ JJ - X_and_dX = ForwardDiff.Dual.(X, ((0, 0),)) - for return_byproduct in (true, false) - res_and_dres = @inferred implicit(X_and_dX, Val(return_byproduct)) - if return_byproduct - Y_and_dY, Z = res_and_dres - @test size(Y_and_dY) == size(Y) - else - Y_and_dY = res_and_dres - @test size(Y_and_dY) == size(Y) - end - end - end - - @testset "Reverse" begin - @test Zygote.jacobian(implicit, X)[1] ≈ JJ - for return_byproduct in (true, false) - _, pullback = @inferred rrule( - Zygote.ZygoteRuleConfig(), implicit, X, Val(return_byproduct) - ) - dY, dZ = zero(implicit(X)), 0 - if return_byproduct - @test (@inferred pullback((dY, dZ))) == pullback((dY, dZ)) - _, dX = pullback((dY, dZ)) - @test size(dX) == size(X) - else - @test (@inferred pullback(dY)) == pullback(dY) - _, dX = pullback(dY) - @test size(dX) == size(X) - end - end - end - end - end -end diff --git a/test/runtests.jl b/test/runtests.jl index d26134b..6918af3 100644 --- a/test/runtests.jl +++ b/test/runtests.jl @@ -38,26 +38,29 @@ EXAMPLES_DIR_JL = joinpath(dirname(@__DIR__), "examples") Aqua.test_undefined_exports(ImplicitDifferentiation) Aqua.test_piracy(ImplicitDifferentiation) Aqua.test_project_extras(ImplicitDifferentiation) - Aqua.test_stale_deps(ImplicitDifferentiation; ignore=[:ChainRulesCore]) + Aqua.test_stale_deps( + ImplicitDifferentiation; ignore=[:AbstractDifferentiation, :ChainRulesCore] + ) Aqua.test_deps_compat(ImplicitDifferentiation) if VERSION >= v"1.7" Aqua.test_project_toml_formatting(ImplicitDifferentiation) end end @testset verbose = true "Formatting (JuliaFormatter.jl)" begin - @test format(ImplicitDifferentiation; verbose=true, overwrite=false) + @test format(ImplicitDifferentiation; verbose=false, overwrite=false) end @testset verbose = true "Static checking (JET.jl)" begin - if VERSION >= v"1.8" - JET.test_package(ImplicitDifferentiation; toplevel_logger=nothing) + if VERSION >= v"1.9" + JET.test_package( + ImplicitDifferentiation; + target_modules=(ImplicitDifferentiation,), + toplevel_logger=nothing, + ) end end @testset verbose = false "Doctests (Documenter.jl)" begin doctest(ImplicitDifferentiation) end - @testset verbose = true "Miscellaneous" begin - include("misc.jl") - end @testset verbose = true "Examples" begin for file in readdir(EXAMPLES_DIR_JL) path = joinpath(EXAMPLES_DIR_JL, file) @@ -67,4 +70,7 @@ EXAMPLES_DIR_JL = joinpath(dirname(@__DIR__), "examples") end end end + @testset verbose = true "Systematic" begin + include("systematic.jl") + end end diff --git a/test/systematic.jl b/test/systematic.jl new file mode 100644 index 0000000..7551962 --- /dev/null +++ b/test/systematic.jl @@ -0,0 +1,76 @@ +using ChainRulesCore +using ChainRulesTestUtils +using ForwardDiff +using ImplicitDifferentiation +using JET +using LinearAlgebra +using Random +using StaticArrays +using Test +using Zygote + +Random.seed!(63); + +""" + mysqrt(x) + +Compute the elementwise square root, breaking Zygote.jl and ForwardDiff.jl in the process. +""" +function mysqrt(x::AbstractArray) + a = [0.0] + a[1] = first(x) + return sqrt.(x) +end + +function make_implicit_sqrt(linear_solver) + forward(x) = mysqrt(x) + conditions(x, y) = y .^ 2 .- x + implicit = ImplicitFunction(forward, conditions, linear_solver) + return implicit +end + +function test_implicit_sqrt_call(implicit, x) + @test (@inferred implicit(x)) ≈ sqrt.(x) + if VERSION >= v"1.9" + test_opt(implicit, (typeof(x),)) + end +end + +function test_implicit_sqrt_forward(implicit, x) + y = implicit(x) + J = Diagonal(0.5 ./ vec(sqrt.(x))) + @test ForwardDiff.jacobian(implicit, x) ≈ J + x_and_dx = ForwardDiff.Dual.(x, ((0, 0),)) + res_and_dres = @inferred implicit(x_and_dx) + y_and_dy = res_and_dres + @test size(y_and_dy) == size(y) +end + +function test_implicit_sqrt_reverse(implicit, x) + J = Diagonal(0.5 ./ vec(sqrt.(x))) + @test Zygote.jacobian(implicit, x)[1] ≈ J + _, pullback = @inferred rrule(Zygote.ZygoteRuleConfig(), implicit, x) + dy = zero(implicit(x)) + @test (@inferred pullback(dy)) == pullback(dy) + _, dx = pullback(dy) + @test size(dx) == size(x) + # Skipped because of https://github.com/JuliaDiff/ChainRulesTestUtils.jl/issues/232 and because it detects weird type instabilities + @test_skip test_rrule(implicit, x) +end + +for x in (rand(2), rand(2, 3, 4), SVector{2}(rand(2)), SArray{Tuple{2,3,4}}(rand(2, 3, 4))), + linear_solver in (IterativeLinearSolver(), DirectLinearSolver()) + + testsetname = "$(typeof(x)) - $(typeof(linear_solver))" + @testset verbose = true "$testsetname" begin + if x isa StaticArray && linear_solver isa IterativeLinearSolver + continue + end + forward(x) = mysqrt(x) + conditions(x, y) = y .^ 2 .- x + implicit = ImplicitFunction(forward, conditions, linear_solver) + test_implicit_sqrt_call(implicit, x) + test_implicit_sqrt_forward(implicit, x) + test_implicit_sqrt_reverse(implicit, x) + end +end From 63923a72497806b0309529d34cbc6fbb212b4967 Mon Sep 17 00:00:00 2001 From: Guillaume Dalle <22795598+gdalle@users.noreply.github.com> Date: Fri, 28 Jul 2023 11:57:46 +0200 Subject: [PATCH 3/8] Fix and better test staticarrays --- docs/make.jl | 11 ++++- docs/src/faq.md | 5 ++- examples/6_byproduct.jl | 8 ---- ext/ImplicitDifferentiationChainRulesExt.jl | 9 ++-- ext/ImplicitDifferentiationForwardDiffExt.jl | 5 +-- ext/ImplicitDifferentiationStaticArraysExt.jl | 10 +---- src/ImplicitDifferentiation.jl | 2 +- src/implicit_function.jl | 6 +-- src/linear_solver.jl | 19 ++++---- test/runtests.jl | 9 ++-- test/systematic.jl | 45 ++++++++++++------- 11 files changed, 70 insertions(+), 59 deletions(-) delete mode 100644 examples/6_byproduct.jl diff --git a/docs/make.jl b/docs/make.jl index 44850de..f58f9b7 100644 --- a/docs/make.jl +++ b/docs/make.jl @@ -48,10 +48,19 @@ for file in sort(readdir(EXAMPLES_DIR_MD)) end end +last_basic_example = 4 +basic_example_pages = filter( + p -> parse(Int, last(p)[1]) <= last_basic_example, example_pages +) +advanced_example_pages = filter( + p -> parse(Int, last(p)[1]) > last_basic_example, example_pages +) + pages = [ "Home" => "index.md", + "Basic examples" => basic_example_pages, + "Advanced examples" => advanced_example_pages, "API reference" => "api.md", - "Examples" => example_pages, "FAQ" => "faq.md", ] diff --git a/docs/src/faq.md b/docs/src/faq.md index 425291f..d92d238 100644 --- a/docs/src/faq.md +++ b/docs/src/faq.md @@ -5,7 +5,10 @@ - Forward mode: [ForwardDiff.jl](https://github.com/JuliaDiff/ForwardDiff.jl) - Reverse mode: all the packages compatible with [ChainRules.jl](https://github.com/JuliaDiff/ChainRules.jl) -In the future, we would like to add [Enzyme.jl](https://github.com/EnzymeAD/Enzyme.jl) support. +In the future, we would like to add + +- [Enzyme.jl](https://github.com/EnzymeAD/Enzyme.jl) +- [Diffractor.jl](https://github.com/JuliaDiff/Diffractor.jl) ## Higher-dimensional arrays diff --git a/examples/6_byproduct.jl b/examples/6_byproduct.jl deleted file mode 100644 index 69f75b7..0000000 --- a/examples/6_byproduct.jl +++ /dev/null @@ -1,8 +0,0 @@ -# # Byproduct - -#= -In this example, we show that returning a byproduct with the forward mapping can be useful in certain situations. - -!!! warning "Work in progress" - Come back soon! -=# diff --git a/ext/ImplicitDifferentiationChainRulesExt.jl b/ext/ImplicitDifferentiationChainRulesExt.jl index d2af6d1..4ee48b4 100644 --- a/ext/ImplicitDifferentiationChainRulesExt.jl +++ b/ext/ImplicitDifferentiationChainRulesExt.jl @@ -3,7 +3,7 @@ module ImplicitDifferentiationChainRulesExt using AbstractDifferentiation: ReverseRuleConfigBackend, pullback_function using ChainRulesCore: ChainRulesCore, NoTangent, RuleConfig, ZeroTangent, rrule, unthunk using ImplicitDifferentiation: - ImplicitFunction, PullbackMul!, ReturnByproduct, check_solution, presolve, solve + ImplicitFunction, PullbackMul!, ReturnByproduct, presolve, solve using LinearAlgebra: lmul!, mul! using LinearOperators: LinearOperator using SimpleUnPack: @unpack @@ -14,7 +14,7 @@ using SimpleUnPack: @unpack Custom reverse rule for an [`ImplicitFunction`](@ref), to ensure compatibility with reverse mode autodiff. -This is only available if ChainRulesCore.jl is loaded (extension). +This is only available if ChainRulesCore.jl is loaded (extension), except on Julia < 1.9 where it is always available. - By default, this returns a single output `y(x)` with a pullback accepting a single cotangent `dy`. - If `ReturnByproduct()` is passed as an argument, this returns a couple of outputs `(y(x),z(x))` with a pullback accepting a couple of cotangents `(dy, dz)` (remember that `z(x)` is not differentiated so its cotangent is ignored). @@ -68,13 +68,12 @@ function (implicit_pullback::ImplicitPullback)((dy, dz)) @unpack Aᵀ_op, Bᵀ_op, linear_solver, x = implicit_pullback R = eltype(x) dy_vec = convert(AbstractVector{R}, vec(unthunk(dy))) - dF_vec, stats = solve(linear_solver, Aᵀ_op, dy_vec) - check_solution(linear_solver, stats) + dF_vec = solve(linear_solver, Aᵀ_op, dy_vec) dx_vec = vec(similar(x)) mul!(dx_vec, Bᵀ_op, dF_vec) lmul!(-one(R), dx_vec) dx = reshape(dx_vec, size(x)) - return (NoTangent(), dx, NoTangent()) + return (NoTangent(), dx) end end diff --git a/ext/ImplicitDifferentiationForwardDiffExt.jl b/ext/ImplicitDifferentiationForwardDiffExt.jl index 8c16f39..2bf89b0 100644 --- a/ext/ImplicitDifferentiationForwardDiffExt.jl +++ b/ext/ImplicitDifferentiationForwardDiffExt.jl @@ -9,7 +9,7 @@ end using AbstractDifferentiation: AbstractDifferentiation, ForwardDiffBackend, pushforward_function using ImplicitDifferentiation: - ImplicitFunction, PushforwardMul!, ReturnByproduct, check_solution, presolve, solve + ImplicitFunction, PushforwardMul!, ReturnByproduct, presolve, solve using LinearAlgebra: lmul!, mul! using LinearOperators: LinearOperator using SimpleUnPack: @unpack @@ -49,9 +49,8 @@ function (implicit::ImplicitFunction)( dₖx_vec = vec(partials.(x_and_dx, k)) Bdx = vec(similar(y)) mul!(Bdx, B_op, dₖx_vec) - dₖy_vec, stats = solve(linear_solver, A_op_presolved, Bdx) + dₖy_vec = solve(linear_solver, A_op_presolved, Bdx) lmul!(-one(R), dₖy_vec) - check_solution(linear_solver, stats) reshape(dₖy_vec, size(y)) end diff --git a/ext/ImplicitDifferentiationStaticArraysExt.jl b/ext/ImplicitDifferentiationStaticArraysExt.jl index 629be7c..bfae8d5 100644 --- a/ext/ImplicitDifferentiationStaticArraysExt.jl +++ b/ext/ImplicitDifferentiationStaticArraysExt.jl @@ -9,17 +9,11 @@ end import ImplicitDifferentiation: ImplicitDifferentiation, DirectLinearSolver using LinearAlgebra: lu, mul! -_prodsize(::Type{Tuple{}}) = 1 -_prodsize(::Type{Tuple{N1}}) where {N1} = N1 -_prodsize(::Type{Tuple{N1,N2}}) where {N1,N2} = N1 * N2 -_prodsize(::Type{Tuple{N1,N2,N3}}) where {N1,N2,N3} = N1 * N2 * N3 -_prodsize(::Type{Tuple{N1,N2,N3,N4}}) where {N1,N2,N3,N4} = N1 * N2 * N3 * N4 - function ImplicitDifferentiation.presolve( ::DirectLinearSolver, A, y::StaticArray{S,T,N} ) where {S,T,N} - S_prod = _prodsize(S) - A_static = zero(MMatrix{S_prod,S_prod,T}) + m = length(y) + A_static = zero(MMatrix{m,m,T}) for i in axes(A_static, 2) v = vec(similar(y)) v .= zero(T) diff --git a/src/ImplicitDifferentiation.jl b/src/ImplicitDifferentiation.jl index dab8038..3b3c446 100644 --- a/src/ImplicitDifferentiation.jl +++ b/src/ImplicitDifferentiation.jl @@ -2,7 +2,7 @@ module ImplicitDifferentiation using Krylov: KrylovStats, gmres using LinearOperators: LinearOperators, LinearOperator -using LinearAlgebra: lu +using LinearAlgebra: lu, SingularException using Requires: @require using SimpleUnPack: @unpack diff --git a/src/implicit_function.jl b/src/implicit_function.jl index 4840007..62d2a89 100644 --- a/src/implicit_function.jl +++ b/src/implicit_function.jl @@ -16,7 +16,7 @@ While `f` does not not need to be compatible with automatic differentiation, `c` # Details - By default, we assume that the forward mapping is `x -> y(x)` and the conditions are `c(x,y(x)) = 0`. -- If `handleByproduct()` is passed as an argument to the constructor, we assume instead that the forward mapping is `x -> (y(x),z(x))` and the conditions are `c(x,y(x),z(x)) = 0`. In this case, `z(x)` can contain additional information generated by the forward mapping, but beware that we consider it constant for differentiation purposes. +- If `HandleByproduct()` is passed as an argument to the constructor, we assume instead that the forward mapping is `x -> (y(x),z(x))` and the conditions are `c(x,y(x),z(x)) = 0`. In this case, `z(x)` can contain additional information generated by the forward mapping, but beware that we consider it constant for differentiation purposes. Given `x ∈ ℝⁿ` and `y ∈ ℝᵈ`, we need as many conditions as output dimensions: `c(x,y,z) ∈ ℝᵈ`. We can then compute the Jacobian of `y(⋅)` using the implicit function theorem: ``` @@ -72,8 +72,8 @@ end Make an [`ImplicitFunction`](@ref) callable by applying the forward mapping `implicit.forward`. """ -function (implicit::ImplicitFunction)(x::AbstractArray, ; kwargs...) - y, z = implicit.forward(x, ; kwargs...) +function (implicit::ImplicitFunction)(x::AbstractArray; kwargs...) + y, z = implicit.forward(x; kwargs...) return y end diff --git a/src/linear_solver.jl b/src/linear_solver.jl index 06a813f..4d9b86d 100644 --- a/src/linear_solver.jl +++ b/src/linear_solver.jl @@ -18,7 +18,14 @@ An implementation of `AbstractLinearSolver` using `Krylov.gmres`. struct IterativeLinearSolver <: AbstractLinearSolver end presolve(::IterativeLinearSolver, A, y) = A -solve(::IterativeLinearSolver, A, b) = gmres(A, b) + +function solve(::IterativeLinearSolver, A, b) + x, stats = gmres(A, b) + if !stats.solved + throw(SolverFailureException(gmres, stats)) + end + return x +end """ DirectLinearSolver @@ -28,7 +35,7 @@ An implementation of `AbstractLinearSolver` using the built-in `\` operator. struct DirectLinearSolver <: AbstractLinearSolver end presolve(::DirectLinearSolver, A, y) = lu(Matrix(A)) -solve(::DirectLinearSolver, A, b) = (A \ b, (; solved=true)) +solve(::DirectLinearSolver, A, b) = A \ b struct SolverFailureException{A,B} <: Exception solver::A @@ -41,11 +48,3 @@ function Base.show(io::IO, sfe::SolverFailureException) "SolverFailureException: \n Linear solver: $(sfe.solver) \n Solver stats: $(string(sfe.stats))", ) end - -function check_solution(solver, stats) - if stats.solved - return nothing - else - throw(SolverFailureException(solver, stats)) - end -end diff --git a/test/runtests.jl b/test/runtests.jl index 6918af3..1256129 100644 --- a/test/runtests.jl +++ b/test/runtests.jl @@ -2,6 +2,7 @@ using Aqua using Documenter +using ForwardDiff using ImplicitDifferentiation using JET using JuliaFormatter @@ -53,7 +54,7 @@ EXAMPLES_DIR_JL = joinpath(dirname(@__DIR__), "examples") if VERSION >= v"1.9" JET.test_package( ImplicitDifferentiation; - target_modules=(ImplicitDifferentiation,), + target_defined_modules=true, toplevel_logger=nothing, ) end @@ -61,6 +62,9 @@ EXAMPLES_DIR_JL = joinpath(dirname(@__DIR__), "examples") @testset verbose = false "Doctests (Documenter.jl)" begin doctest(ImplicitDifferentiation) end + @testset verbose = true "Systematic" begin + include("systematic.jl") + end @testset verbose = true "Examples" begin for file in readdir(EXAMPLES_DIR_JL) path = joinpath(EXAMPLES_DIR_JL, file) @@ -70,7 +74,4 @@ EXAMPLES_DIR_JL = joinpath(dirname(@__DIR__), "examples") end end end - @testset verbose = true "Systematic" begin - include("systematic.jl") - end end diff --git a/test/systematic.jl b/test/systematic.jl index 7551962..bb5fce1 100644 --- a/test/systematic.jl +++ b/test/systematic.jl @@ -31,6 +31,9 @@ end function test_implicit_sqrt_call(implicit, x) @test (@inferred implicit(x)) ≈ sqrt.(x) + if typeof(x) <: StaticArray + @test typeof(implicit(x)) <: StaticArray + end if VERSION >= v"1.9" test_opt(implicit, (typeof(x),)) end @@ -41,36 +44,48 @@ function test_implicit_sqrt_forward(implicit, x) J = Diagonal(0.5 ./ vec(sqrt.(x))) @test ForwardDiff.jacobian(implicit, x) ≈ J x_and_dx = ForwardDiff.Dual.(x, ((0, 0),)) - res_and_dres = @inferred implicit(x_and_dx) - y_and_dy = res_and_dres + y_and_dy = @inferred implicit(x_and_dx) @test size(y_and_dy) == size(y) + if typeof(x) <: StaticArray + @test typeof(y_and_dy) <: StaticArray || + typeof(y_and_dy) <: (Base.ReshapedArray{T,N,<:StaticArray} where {T,N}) + end end function test_implicit_sqrt_reverse(implicit, x) J = Diagonal(0.5 ./ vec(sqrt.(x))) @test Zygote.jacobian(implicit, x)[1] ≈ J _, pullback = @inferred rrule(Zygote.ZygoteRuleConfig(), implicit, x) - dy = zero(implicit(x)) + dy = similar(implicit(x)) + dy .= zero(eltype(dy)) @test (@inferred pullback(dy)) == pullback(dy) _, dx = pullback(dy) @test size(dx) == size(x) + if typeof(x) <: StaticArray + @test typeof(dx) <: StaticArray || + typeof(dx) <: (Base.ReshapedArray{T,N,<:StaticArray} where {T,N}) + end # Skipped because of https://github.com/JuliaDiff/ChainRulesTestUtils.jl/issues/232 and because it detects weird type instabilities @test_skip test_rrule(implicit, x) end -for x in (rand(2), rand(2, 3, 4), SVector{2}(rand(2)), SArray{Tuple{2,3,4}}(rand(2, 3, 4))), - linear_solver in (IterativeLinearSolver(), DirectLinearSolver()) +x_candidates = ( + rand(2), rand(2, 3, 4), SVector{2}(rand(2)), SArray{Tuple{2,3,4}}(rand(2, 3, 4)) +) - testsetname = "$(typeof(x)) - $(typeof(linear_solver))" - @testset verbose = true "$testsetname" begin - if x isa StaticArray && linear_solver isa IterativeLinearSolver - continue +for linear_solver in (IterativeLinearSolver(), DirectLinearSolver()) + forward(x) = mysqrt(x) + conditions(x, y) = y .^ 2 .- x + implicit = ImplicitFunction(forward, conditions, linear_solver) + for x in x_candidates + testsetname = "$(typeof(x)) - $(typeof(linear_solver))" + @testset verbose = true "$testsetname" begin + if x isa StaticArray && linear_solver isa IterativeLinearSolver + continue + end + test_implicit_sqrt_call(implicit, x) + test_implicit_sqrt_forward(implicit, x) + test_implicit_sqrt_reverse(implicit, x) end - forward(x) = mysqrt(x) - conditions(x, y) = y .^ 2 .- x - implicit = ImplicitFunction(forward, conditions, linear_solver) - test_implicit_sqrt_call(implicit, x) - test_implicit_sqrt_forward(implicit, x) - test_implicit_sqrt_reverse(implicit, x) end end From a278a4faab206eb0d7622f358aad18a84de7f5d0 Mon Sep 17 00:00:00 2001 From: Guillaume Dalle <22795598+gdalle@users.noreply.github.com> Date: Fri, 28 Jul 2023 17:34:16 +0200 Subject: [PATCH 4/8] Add application for byproducts & group docs examples --- docs/Manifest.toml | 209 ++++++++------- docs/make.jl | 15 +- examples/{0_basic.jl => 0_intro.jl} | 4 +- examples/1_basic.jl | 242 ++++++++++++++++++ examples/1_unconstrained_optim.jl | 99 ------- .../{4_constrained_optim.jl => 2_advanced.jl} | 31 ++- examples/2_nonlinear_solve.jl | 86 ------- examples/3_fixed_points.jl | 85 ------ examples/{5_multiargs.jl => 3_tricks.jl} | 86 ++++++- 9 files changed, 459 insertions(+), 398 deletions(-) rename examples/{0_basic.jl => 0_intro.jl} (97%) create mode 100644 examples/1_basic.jl delete mode 100644 examples/1_unconstrained_optim.jl rename examples/{4_constrained_optim.jl => 2_advanced.jl} (88%) delete mode 100644 examples/2_nonlinear_solve.jl delete mode 100644 examples/3_fixed_points.jl rename examples/{5_multiargs.jl => 3_tricks.jl} (51%) diff --git a/docs/Manifest.toml b/docs/Manifest.toml index bdecaf8..9b9c6fd 100644 --- a/docs/Manifest.toml +++ b/docs/Manifest.toml @@ -1,6 +1,6 @@ # This file is machine-generated - editing it directly is not advised -julia_version = "1.9.0" +julia_version = "1.9.2" manifest_format = "2.0" project_hash = "e7f4896b7e8c3921c7466749f467ab7680867992" @@ -40,9 +40,9 @@ version = "0.5.2" [[deps.AbstractFFTs]] deps = ["LinearAlgebra"] -git-tree-sha1 = "16b6dbc4cf7caee4e1e75c49485ec67b667098a0" +git-tree-sha1 = "cad4c758c0038eea30394b1b671526921ca85b21" uuid = "621f4979-c628-5d54-868e-fcf4e3e8185c" -version = "1.3.1" +version = "1.4.0" weakdeps = ["ChainRulesCore"] [deps.AbstractFFTs.extensions] @@ -50,9 +50,9 @@ weakdeps = ["ChainRulesCore"] [[deps.Adapt]] deps = ["LinearAlgebra", "Requires"] -git-tree-sha1 = "cc37d689f599e8df4f464b2fa3870ff7db7492ef" +git-tree-sha1 = "76289dc51920fdc6e0013c872ba9551d54961c24" uuid = "79e6a3ab-5dfb-504d-930d-738a2a938a0e" -version = "3.6.1" +version = "3.6.2" weakdeps = ["StaticArrays"] [deps.Adapt.extensions] @@ -63,10 +63,10 @@ uuid = "0dad84c5-d112-42e6-8d28-ef12dabb789f" version = "1.1.1" [[deps.ArrayInterface]] -deps = ["Adapt", "LinearAlgebra", "Requires", "SnoopPrecompile", "SparseArrays", "SuiteSparse"] -git-tree-sha1 = "38911c7737e123b28182d89027f4216cfc8a9da7" +deps = ["Adapt", "LinearAlgebra", "Requires", "SparseArrays", "SuiteSparse"] +git-tree-sha1 = "f83ec24f76d4c8f525099b2ac475fc098138ec31" uuid = "4fba245c-0d91-5ea0-9b3e-6abc04ee57a9" -version = "7.4.3" +version = "7.4.11" [deps.ArrayInterface.extensions] ArrayInterfaceBandedMatricesExt = "BandedMatrices" @@ -97,21 +97,21 @@ version = "0.4.2" [[deps.ChainRules]] deps = ["Adapt", "ChainRulesCore", "Compat", "Distributed", "GPUArraysCore", "IrrationalConstants", "LinearAlgebra", "Random", "RealDot", "SparseArrays", "Statistics", "StructArrays"] -git-tree-sha1 = "7d20c2fb8ab838e41069398685e7b6b5f89ed85b" +git-tree-sha1 = "f98ae934cd677d51d2941088849f0bf2f59e6f6e" uuid = "082447d4-558c-5d27-93f4-14fc19e9eca2" -version = "1.48.0" +version = "1.53.0" [[deps.ChainRulesCore]] deps = ["Compat", "LinearAlgebra", "SparseArrays"] -git-tree-sha1 = "c6d890a52d2c4d55d326439580c3b8d0875a77d9" +git-tree-sha1 = "e30f2f4e20f7f186dc36529910beaedc60cfa644" uuid = "d360d2e6-b24c-11e9-a2a3-2a2ae2dbcce4" -version = "1.15.7" +version = "1.16.0" [[deps.ChainRulesTestUtils]] deps = ["ChainRulesCore", "Compat", "FiniteDifferences", "LinearAlgebra", "Random", "Test"] -git-tree-sha1 = "1237bdbcfec728721718ef57dcb855a19c11bf3a" +git-tree-sha1 = "5ab2a7bc21ecc3eb0226478ff8f87e9685b11818" uuid = "cdddcdb0-9152-4a09-a978-84456f9df70a" -version = "1.10.1" +version = "1.11.0" [[deps.CommonSubexpressions]] deps = ["MacroTools", "Test"] @@ -121,9 +121,9 @@ version = "0.3.0" [[deps.Compat]] deps = ["UUIDs"] -git-tree-sha1 = "7a60c856b9fa189eb34f5f8a6f6b5529b7942957" +git-tree-sha1 = "5ce999a19f4ca23ea484e92a1774a61b8ca4cf8e" uuid = "34da2185-b29b-5c13-b0c7-acf172513d20" -version = "4.6.1" +version = "4.8.0" weakdeps = ["Dates", "LinearAlgebra"] [deps.Compat.extensions] @@ -132,13 +132,13 @@ weakdeps = ["Dates", "LinearAlgebra"] [[deps.CompilerSupportLibraries_jll]] deps = ["Artifacts", "Libdl"] uuid = "e66e0078-7015-5450-92f7-15fbd957f2ae" -version = "1.0.2+0" +version = "1.0.5+0" [[deps.ComponentArrays]] deps = ["ArrayInterface", "ChainRulesCore", "ForwardDiff", "Functors", "LinearAlgebra", "Requires", "StaticArrayInterface"] -git-tree-sha1 = "891f08177789faff56f0deda1e23615ec220ce44" +git-tree-sha1 = "e1a6694a62f7a00cf2c4f65dccd118ac28a6c099" uuid = "b0b7db55-cfe3-40fc-9ded-d10e2dbeff66" -version = "0.13.12" +version = "0.14.0" [deps.ComponentArrays.extensions] ComponentArraysConstructionBaseExt = "ConstructionBase" @@ -158,28 +158,28 @@ version = "0.13.12" [[deps.ConstructionBase]] deps = ["LinearAlgebra"] -git-tree-sha1 = "89a9db8d28102b094992472d333674bd1a83ce2a" +git-tree-sha1 = "fe2838a593b5f776e1597e086dcd47560d94e816" uuid = "187b0558-2788-49d3-abe0-74a17ed4e7c9" -version = "1.5.1" +version = "1.5.3" [deps.ConstructionBase.extensions] - IntervalSetsExt = "IntervalSets" - StaticArraysExt = "StaticArrays" + ConstructionBaseIntervalSetsExt = "IntervalSets" + ConstructionBaseStaticArraysExt = "StaticArrays" [deps.ConstructionBase.weakdeps] IntervalSets = "8197267c-284f-5f27-9208-e0e47529a953" StaticArrays = "90137ffa-7385-5640-81b9-e52037218182" [[deps.DataAPI]] -git-tree-sha1 = "e8119c1a33d267e16108be441a287a6981ba1630" +git-tree-sha1 = "8da84edb865b0b5b0100c0666a9bc9a0b71c553c" uuid = "9a962f9c-6df0-11e9-0e5d-c546b8b5ee8a" -version = "1.14.0" +version = "1.15.0" [[deps.DataStructures]] deps = ["Compat", "InteractiveUtils", "OrderedCollections"] -git-tree-sha1 = "d1fff3a548102f48987a52a2e0d114fa97d730f0" +git-tree-sha1 = "cf25ccb972fec4e4817764d01c82386ae94f77b4" uuid = "864edb3b-99cc-5e75-8d2d-829cb0a9cfe8" -version = "0.18.13" +version = "0.18.14" [[deps.DataValueInterfaces]] git-tree-sha1 = "bfc1187b79289637fa0ef6d4436ebdfe6905cbd6" @@ -198,15 +198,19 @@ version = "1.1.0" [[deps.DiffRules]] deps = ["IrrationalConstants", "LogExpFunctions", "NaNMath", "Random", "SpecialFunctions"] -git-tree-sha1 = "a4ad7ef19d2cdc2eff57abbbe68032b1cd0bd8f8" +git-tree-sha1 = "23163d55f885173722d1e4cf0f6110cdbaf7e272" uuid = "b552c78f-8df3-52c6-915a-8e097449b14b" -version = "1.13.0" +version = "1.15.1" [[deps.Distances]] -deps = ["LinearAlgebra", "SparseArrays", "Statistics", "StatsAPI"] -git-tree-sha1 = "49eba9ad9f7ead780bfb7ee319f962c811c6d3b2" +deps = ["LinearAlgebra", "Statistics", "StatsAPI"] +git-tree-sha1 = "b6def76ffad15143924a2199f72a5cd883a2e8a9" uuid = "b4f34e82-e78d-54a5-968a-f98e89d6e8f7" -version = "0.10.8" +version = "0.10.9" +weakdeps = ["SparseArrays"] + + [deps.Distances.extensions] + DistancesSparseArraysExt = "SparseArrays" [[deps.Distributed]] deps = ["Random", "Serialization", "Sockets"] @@ -220,9 +224,9 @@ version = "0.9.3" [[deps.Documenter]] deps = ["ANSIColoredPrinters", "Base64", "Dates", "DocStringExtensions", "IOCapture", "InteractiveUtils", "JSON", "LibGit2", "Logging", "Markdown", "REPL", "Test", "Unicode"] -git-tree-sha1 = "58fea7c536acd71f3eef6be3b21c0df5f3df88fd" +git-tree-sha1 = "39fd748a73dce4c05a9655475e437170d8fb1b67" uuid = "e30172f5-a6a5-5a46-863b-614d45cd2de4" -version = "0.27.24" +version = "0.27.25" [[deps.Downloads]] deps = ["ArgTools", "FileWatching", "LibCURL", "NetworkOptions"] @@ -230,9 +234,9 @@ uuid = "f43a241f-c20a-4ad4-852c-f6b1247861c6" version = "1.6.0" [[deps.ExprTools]] -git-tree-sha1 = "c1d06d129da9f55715c6c212866f5b1bddc5fa00" +git-tree-sha1 = "27415f162e6028e81c72b82ef756bf321213b6ec" uuid = "e2ba6199-217a-4e67-a87a-7c52f15ade04" -version = "0.1.9" +version = "0.1.10" [[deps.FastClosures]] git-tree-sha1 = "acebe244d53ee1b461970f8910c235b259e772ef" @@ -244,21 +248,31 @@ uuid = "7b1f6079-737a-58dc-b8bc-7a2ca5c1b5ee" [[deps.FillArrays]] deps = ["LinearAlgebra", "Random", "SparseArrays", "Statistics"] -git-tree-sha1 = "7072f1e3e5a8be51d525d64f63d3ec1287ff2790" +git-tree-sha1 = "f372472e8672b1d993e93dada09e23139b509f9e" uuid = "1a297f60-69ca-5386-bcde-b61e274b549b" -version = "0.13.11" +version = "1.5.0" [[deps.FiniteDiff]] -deps = ["ArrayInterface", "LinearAlgebra", "Requires", "Setfield", "SparseArrays", "StaticArrays"] -git-tree-sha1 = "03fcb1c42ec905d15b305359603888ec3e65f886" +deps = ["ArrayInterface", "LinearAlgebra", "Requires", "Setfield", "SparseArrays"] +git-tree-sha1 = "c6e4a1fbe73b31a3dea94b1da449503b8830c306" uuid = "6a86dc24-6348-571c-b903-95158fe2bd41" -version = "2.19.0" +version = "2.21.1" + + [deps.FiniteDiff.extensions] + FiniteDiffBandedMatricesExt = "BandedMatrices" + FiniteDiffBlockBandedMatricesExt = "BlockBandedMatrices" + FiniteDiffStaticArraysExt = "StaticArrays" + + [deps.FiniteDiff.weakdeps] + BandedMatrices = "aae01518-5342-5314-be14-df237901396f" + BlockBandedMatrices = "ffab5731-97b5-5995-9138-79e8c1846df0" + StaticArrays = "90137ffa-7385-5640-81b9-e52037218182" [[deps.FiniteDifferences]] deps = ["ChainRulesCore", "LinearAlgebra", "Printf", "Random", "Richardson", "SparseArrays", "StaticArrays"] -git-tree-sha1 = "3f605dd6db5640c5278f2551afc9427656439f42" +git-tree-sha1 = "549636fd813ddf1816d8407efb23f486822f4b63" uuid = "26cc04aa-876d-5657-8c51-4c34ba976000" -version = "0.12.26" +version = "0.12.29" [[deps.ForwardDiff]] deps = ["CommonSubexpressions", "DiffResults", "DiffRules", "LinearAlgebra", "LogExpFunctions", "NaNMath", "Preferences", "Printf", "Random", "SpecialFunctions"] @@ -272,9 +286,9 @@ weakdeps = ["StaticArrays"] [[deps.Functors]] deps = ["LinearAlgebra"] -git-tree-sha1 = "478f8c3145bb91d82c2cf20433e8c1b30df454cc" +git-tree-sha1 = "9a68d75d466ccc1218d0552a8e1631151c569545" uuid = "d9f16b24-f501-4c13-a1f2-28368ffc5196" -version = "0.4.4" +version = "0.4.5" [[deps.Future]] deps = ["Random"] @@ -282,27 +296,27 @@ uuid = "9fa8497b-333b-5362-9e8d-4d0656e87820" [[deps.GPUArrays]] deps = ["Adapt", "GPUArraysCore", "LLVM", "LinearAlgebra", "Printf", "Random", "Reexport", "Serialization", "Statistics"] -git-tree-sha1 = "9ade6983c3dbbd492cf5729f865fe030d1541463" +git-tree-sha1 = "2e57b4a4f9cc15e85a24d603256fe08e527f48d1" uuid = "0c68f7d7-f131-5f86-a1c3-88cf8149b2d7" -version = "8.6.6" +version = "8.8.1" [[deps.GPUArraysCore]] deps = ["Adapt"] -git-tree-sha1 = "1cd7f0af1aa58abc02ea1d872953a97359cb87fa" +git-tree-sha1 = "2d6ca471a6c7b536127afccfa7564b5b39227fe0" uuid = "46192b85-c4d5-4398-a991-12ede77f4527" -version = "0.1.4" +version = "0.1.5" [[deps.IOCapture]] deps = ["Logging", "Random"] -git-tree-sha1 = "f7be53659ab06ddc986428d3a9dcc95f6fa6705a" +git-tree-sha1 = "d75853a0bdbfb1ac815478bacd89cd27b550ace6" uuid = "b5f81e59-6552-4d32-b1f0-c071b021bf89" -version = "0.2.2" +version = "0.2.3" [[deps.IRTools]] deps = ["InteractiveUtils", "MacroTools", "Test"] -git-tree-sha1 = "0ade27f0c49cebd8db2523c4eeccf779407cf12c" +git-tree-sha1 = "eac00994ce3229a464c2847e956d77a2c64ad3a5" uuid = "7869d1d1-7146-5819-86e3-90919afe41df" -version = "0.4.9" +version = "0.4.10" [[deps.IfElse]] git-tree-sha1 = "debdd00ffef04665ccbb3e150747a77560e8fad1" @@ -310,15 +324,16 @@ uuid = "615f187c-cbe4-4ef1-ba3b-2fcf58d6d173" version = "0.1.1" [[deps.ImplicitDifferentiation]] -deps = ["AbstractDifferentiation", "Krylov", "LinearOperators", "Requires", "SimpleUnPack"] +deps = ["AbstractDifferentiation", "Krylov", "LinearAlgebra", "LinearOperators", "Requires", "SimpleUnPack"] path = ".." uuid = "57b37032-215b-411a-8a7c-41a003a55207" version = "0.5.0-DEV" -weakdeps = ["ChainRulesCore", "ForwardDiff"] +weakdeps = ["ChainRulesCore", "ForwardDiff", "StaticArrays"] [deps.ImplicitDifferentiation.extensions] ImplicitDifferentiationChainRulesExt = "ChainRulesCore" ImplicitDifferentiationForwardDiffExt = "ForwardDiff" + ImplicitDifferentiationStaticArraysExt = "StaticArrays" [[deps.InteractiveUtils]] deps = ["Markdown"] @@ -348,9 +363,9 @@ version = "0.21.4" [[deps.Krylov]] deps = ["LinearAlgebra", "Printf", "SparseArrays"] -git-tree-sha1 = "dd90aacbfb622f898a97c2a4411ac49101ebab8a" +git-tree-sha1 = "6dc4ad9cd74ad4ca0a8e219e945dbd22039f2125" uuid = "ba0b0d4f-ebba-5204-a429-3ac8c609bfb7" -version = "0.9.0" +version = "0.9.2" [[deps.LDLFactorizations]] deps = ["AMD", "LinearAlgebra", "SparseArrays", "Test"] @@ -360,15 +375,15 @@ version = "0.10.0" [[deps.LLVM]] deps = ["CEnum", "LLVMExtra_jll", "Libdl", "Printf", "Unicode"] -git-tree-sha1 = "a8960cae30b42b66dd41808beb76490519f6f9e2" +git-tree-sha1 = "8695a49bfe05a2dc0feeefd06b4ca6361a018729" uuid = "929cbde3-209d-540e-8aea-75f648917ca0" -version = "5.0.0" +version = "6.1.0" [[deps.LLVMExtra_jll]] deps = ["Artifacts", "JLLWrappers", "LazyArtifacts", "Libdl", "TOML"] -git-tree-sha1 = "09b7505cc0b1cee87e5d4a26eea61d2e1b0dcd35" +git-tree-sha1 = "c35203c1e1002747da220ffc3c0762ce7754b08c" uuid = "dad2f222-ce93-54a1-a47d-0025e8a3acab" -version = "0.0.21+0" +version = "0.0.23+0" [[deps.LazyArtifacts]] deps = ["Artifacts", "Pkg"] @@ -420,9 +435,9 @@ version = "2.14.0" [[deps.LogExpFunctions]] deps = ["DocStringExtensions", "IrrationalConstants", "LinearAlgebra"] -git-tree-sha1 = "0a1b7c2863e44523180fdb3146534e265a91870b" +git-tree-sha1 = "c3ce8e7420b3a6e071e0fe4745f5d4300e37b13f" uuid = "2ab3a3ac-af41-5b50-aa03-7779005ae688" -version = "0.3.23" +version = "0.3.24" [deps.LogExpFunctions.extensions] LogExpFunctionsChainRulesCoreExt = "ChainRulesCore" @@ -505,14 +520,14 @@ version = "0.5.5+0" [[deps.Optim]] deps = ["Compat", "FillArrays", "ForwardDiff", "LineSearches", "LinearAlgebra", "NLSolversBase", "NaNMath", "Parameters", "PositiveFactorizations", "Printf", "SparseArrays", "StatsBase"] -git-tree-sha1 = "a89b11f0f354f06099e4001c151dffad7ebab015" +git-tree-sha1 = "e3a6546c1577bfd701771b477b794a52949e7594" uuid = "429524aa-4258-5aef-a3af-852621145aeb" -version = "1.7.5" +version = "1.7.6" [[deps.OrderedCollections]] -git-tree-sha1 = "d321bf2de576bf25ec4d3e4360faca399afca282" +git-tree-sha1 = "2e73fe17cac3c62ad1aebe70d44c963c3cfdc3e3" uuid = "bac558e1-5e72-5ebc-8fee-abe8a469f55d" -version = "1.6.0" +version = "1.6.2" [[deps.Parameters]] deps = ["OrderedCollections", "UnPack"] @@ -521,15 +536,15 @@ uuid = "d96e819e-fc66-5662-9728-84c9c7592b0a" version = "0.12.3" [[deps.Parsers]] -deps = ["Dates", "SnoopPrecompile"] -git-tree-sha1 = "478ac6c952fddd4399e71d4779797c538d0ff2bf" +deps = ["Dates", "PrecompileTools", "UUIDs"] +git-tree-sha1 = "4b2e829ee66d4218e0cef22c0a64ee37cf258c29" uuid = "69de0a69-1ddd-5017-9359-2bf0b02dc9f0" -version = "2.5.8" +version = "2.7.1" [[deps.Pkg]] deps = ["Artifacts", "Dates", "Downloads", "FileWatching", "LibGit2", "Libdl", "Logging", "Markdown", "Printf", "REPL", "Random", "SHA", "Serialization", "TOML", "Tar", "UUIDs", "p7zip_jll"] uuid = "44cfe95a-1eb2-52ea-b672-e2afdf69b78f" -version = "1.9.0" +version = "1.9.2" [[deps.PositiveFactorizations]] deps = ["LinearAlgebra"] @@ -537,11 +552,17 @@ git-tree-sha1 = "17275485f373e6673f7e7f97051f703ed5b15b20" uuid = "85a6dd25-e78a-55b7-8502-1745935b8125" version = "0.2.4" +[[deps.PrecompileTools]] +deps = ["Preferences"] +git-tree-sha1 = "9673d39decc5feece56ef3940e5dafba15ba0f81" +uuid = "aea7be01-6a6a-4083-8856-8a6e6704d82a" +version = "1.1.2" + [[deps.Preferences]] deps = ["TOML"] -git-tree-sha1 = "47e5f437cc0e7ef2ce8406ce1e7e24d44915f88d" +git-tree-sha1 = "7eb1686b4f04b82f96ed7a4ea5890a4f0c7a09f1" uuid = "21216c6a-2e73-6563-6e65-726566657250" -version = "1.3.0" +version = "1.4.0" [[deps.Printf]] deps = ["Unicode"] @@ -607,9 +628,9 @@ uuid = "6462fe0b-24de-5631-8697-dd941f90decc" [[deps.SortingAlgorithms]] deps = ["DataStructures"] -git-tree-sha1 = "a4ada03f999bd01b3a25dcaa30b2d929fe537e00" +git-tree-sha1 = "c60ec5c62180f27efea3ba2908480f8055e17cee" uuid = "a2af1166-a08f-5f64-846c-94a0d3cef48c" -version = "1.1.0" +version = "1.1.1" [[deps.SparseArrays]] deps = ["Libdl", "LinearAlgebra", "Random", "Serialization", "SuiteSparse_jll"] @@ -617,9 +638,9 @@ uuid = "2f01184e-e22b-5df5-ae63-d93ebab69eaf" [[deps.SpecialFunctions]] deps = ["IrrationalConstants", "LogExpFunctions", "OpenLibm_jll", "OpenSpecFun_jll"] -git-tree-sha1 = "ef28127915f4229c971eb43f3fc075dd3fe91880" +git-tree-sha1 = "7beb031cf8145577fbccacd94b8a8f4ce78428d3" uuid = "276daf66-3868-5448-9aa4-cd146d93841b" -version = "2.2.0" +version = "2.3.0" weakdeps = ["ChainRulesCore"] [deps.SpecialFunctions.extensions] @@ -627,9 +648,9 @@ weakdeps = ["ChainRulesCore"] [[deps.Static]] deps = ["IfElse"] -git-tree-sha1 = "dbde6766fc677423598138a5951269432b0fcc90" +git-tree-sha1 = "f295e0a1da4ca425659c57441bcb59abb035a4bc" uuid = "aedffcd0-7271-4cad-89d0-dc628f76c6d3" -version = "0.8.7" +version = "0.8.8" [[deps.StaticArrayInterface]] deps = ["ArrayInterface", "Compat", "IfElse", "LinearAlgebra", "Requires", "SnoopPrecompile", "SparseArrays", "Static", "SuiteSparse"] @@ -646,15 +667,19 @@ version = "1.4.0" StaticArrays = "90137ffa-7385-5640-81b9-e52037218182" [[deps.StaticArrays]] -deps = ["LinearAlgebra", "Random", "StaticArraysCore", "Statistics"] -git-tree-sha1 = "63e84b7fdf5021026d0f17f76af7c57772313d99" +deps = ["LinearAlgebra", "Random", "StaticArraysCore"] +git-tree-sha1 = "9cabadf6e7cd2349b6cf49f1915ad2028d65e881" uuid = "90137ffa-7385-5640-81b9-e52037218182" -version = "1.5.21" +version = "1.6.2" +weakdeps = ["Statistics"] + + [deps.StaticArrays.extensions] + StaticArraysStatisticsExt = "Statistics" [[deps.StaticArraysCore]] -git-tree-sha1 = "6b7ba252635a5eff6a0b0664a41ee140a1c9e72a" +git-tree-sha1 = "36b3d696ce6366023a0ea192b4cd442268995a0d" uuid = "1e83bf80-4336-4d27-bf5d-d5a4f845583c" -version = "1.4.0" +version = "1.4.2" [[deps.Statistics]] deps = ["LinearAlgebra", "SparseArrays"] @@ -669,9 +694,9 @@ version = "1.6.0" [[deps.StatsBase]] deps = ["DataAPI", "DataStructures", "LinearAlgebra", "LogExpFunctions", "Missings", "Printf", "Random", "SortingAlgorithms", "SparseArrays", "Statistics", "StatsAPI"] -git-tree-sha1 = "d1bf48bfcc554a3761a133fe3a9bb01488e06916" +git-tree-sha1 = "75ebe04c5bed70b91614d684259b661c9e6274a4" uuid = "2913bbd2-ae8a-5f71-8c99-4fb6c76f3a91" -version = "0.33.21" +version = "0.34.0" [[deps.StructArrays]] deps = ["Adapt", "DataAPI", "GPUArraysCore", "StaticArraysCore", "Tables"] @@ -716,9 +741,9 @@ uuid = "8dfed614-e22c-5e08-85e1-65c5234f0b40" [[deps.TimerOutputs]] deps = ["ExprTools", "Printf"] -git-tree-sha1 = "f2fd3f288dfc6f507b0c3a2eb3bac009251e548b" +git-tree-sha1 = "f548a9e9c490030e545f72074a41edfd0e5bcdd7" uuid = "a759f4b9-e2f1-59dc-863e-4aeb61b1ea8f" -version = "0.5.22" +version = "0.5.23" [[deps.UUIDs]] deps = ["Random", "SHA"] @@ -738,10 +763,10 @@ uuid = "83775a58-1f1d-513f-b197-d71354ab007a" version = "1.2.13+0" [[deps.Zygote]] -deps = ["AbstractFFTs", "ChainRules", "ChainRulesCore", "DiffRules", "Distributed", "FillArrays", "ForwardDiff", "GPUArrays", "GPUArraysCore", "IRTools", "InteractiveUtils", "LinearAlgebra", "LogExpFunctions", "MacroTools", "NaNMath", "Random", "Requires", "SnoopPrecompile", "SparseArrays", "SpecialFunctions", "Statistics", "ZygoteRules"] -git-tree-sha1 = "987ae5554ca90e837594a0f30325eeb5e7303d1e" +deps = ["AbstractFFTs", "ChainRules", "ChainRulesCore", "DiffRules", "Distributed", "FillArrays", "ForwardDiff", "GPUArrays", "GPUArraysCore", "IRTools", "InteractiveUtils", "LinearAlgebra", "LogExpFunctions", "MacroTools", "NaNMath", "PrecompileTools", "Random", "Requires", "SparseArrays", "SpecialFunctions", "Statistics", "ZygoteRules"] +git-tree-sha1 = "5be3ddb88fc992a7d8ea96c3f10a49a7e98ebc7b" uuid = "e88e6eb3-aa80-5325-afca-941959d7151f" -version = "0.6.60" +version = "0.6.62" [deps.Zygote.extensions] ZygoteColorsExt = "Colors" @@ -762,7 +787,7 @@ version = "0.2.3" [[deps.libblastrampoline_jll]] deps = ["Artifacts", "Libdl"] uuid = "8e850b90-86db-534c-a0d3-1478176c7d93" -version = "5.7.0+0" +version = "5.8.0+0" [[deps.nghttp2_jll]] deps = ["Artifacts", "Libdl"] diff --git a/docs/make.jl b/docs/make.jl index f58f9b7..515c1e1 100644 --- a/docs/make.jl +++ b/docs/make.jl @@ -48,23 +48,14 @@ for file in sort(readdir(EXAMPLES_DIR_MD)) end end -last_basic_example = 4 -basic_example_pages = filter( - p -> parse(Int, last(p)[1]) <= last_basic_example, example_pages -) -advanced_example_pages = filter( - p -> parse(Int, last(p)[1]) > last_basic_example, example_pages -) - pages = [ "Home" => "index.md", - "Basic examples" => basic_example_pages, - "Advanced examples" => advanced_example_pages, + "Examples" => example_pages, "API reference" => "api.md", "FAQ" => "faq.md", ] -format = Documenter.HTML(; +fmt = Documenter.HTML(; prettyurls=get(ENV, "CI", "false") == "true", canonical="https://gdalle.github.io/ImplicitDifferentiation.jl", assets=String[], @@ -76,7 +67,7 @@ makedocs(; authors="Guillaume Dalle, Mohamed Tarek and contributors", repo="https://github.com/gdalle/ImplicitDifferentiation.jl/blob/{commit}{path}#{line}", sitename="ImplicitDifferentiation.jl", - format=format, + format=fmt, pages=pages, linkcheck=true, ) diff --git a/examples/0_basic.jl b/examples/0_intro.jl similarity index 97% rename from examples/0_basic.jl rename to examples/0_intro.jl index 20ca4f9..4e3df7e 100644 --- a/examples/0_basic.jl +++ b/examples/0_intro.jl @@ -1,7 +1,7 @@ -# # Basic use +# # Introduction #= -In this example, we demonstrate the basics of our package on a simple function that is not amenable to automatic differentiation. +We explain the basics of our package on a simple function that is not amenable to naive automatic differentiation. =# using ChainRulesCore #src diff --git a/examples/1_basic.jl b/examples/1_basic.jl new file mode 100644 index 0000000..69b5370 --- /dev/null +++ b/examples/1_basic.jl @@ -0,0 +1,242 @@ +# # Basic use cases + +#= +We show how to differentiate through very common routines: +- an unconstrained optimization problem +- a nonlinear system of equations +- a fixed point iteration +=# + +using ForwardDiff +using ImplicitDifferentiation +using LinearAlgebra +using NLsolve +using Optim +using Random +using Test #src +using Zygote + +Random.seed!(63); + +#= +In all three cases, we will use the square root as our forward mapping, but expressed in three different ways. +Here's our heroic test vector: +=# + +x = rand(2); + +#= +Since we already know the mathematical expression of the Jacobian, we will be able to compare it with our numerical results. +=# + +J = Diagonal(0.5 ./ sqrt.(x)) + +# ## Unconstrained optimization + +#= +First, we show how to differentiate through the solution of an unconstrained optimization problem: +```math +y(x) = \underset{y \in \mathbb{R}^m}{\mathrm{argmin}} ~ f(x, y) +``` +The optimality conditions are given by gradient stationarity: +```math +\nabla_2 f(x, y) = 0 +``` +=# + +#= +To make verification easy, we minimize the following objective: +```math +f(x, y) = \lVert y \odot y - x \rVert^2 +``` +In this case, the optimization problem boils down to the componentwise square root function, but we implement it using a black box solver from [Optim.jl](https://github.com/JuliaNLSolvers/Optim.jl). +Note the presence of a keyword argument. +=# + +function forward_optim(x; method) + f(y) = sum(abs2, y .^ 2 .- x) + y0 = ones(eltype(x), size(x)) + result = optimize(f, y0, method) + return Optim.minimizer(result) +end + +#= +Even though they are defined as a gradient, it is better to provide optimality conditions explicitly: that way we avoid nesting autodiff calls. By default, the conditions should accept two arguments as input. +The forward mapping and the conditions should accept the same set of keyword arguments. +=# + +function conditions_optim(x, y; method) + ∇₂f = 2 .* (y .^ 2 .- x) + return ∇₂f +end + +#= +We now have all the ingredients to construct our implicit function. +=# + +implicit_optim = ImplicitFunction(forward_optim, conditions_optim) + +# And indeed, it behaves as it should when we call it: + +implicit_optim(x; method=LBFGS()) .^ 2 +@test implicit_optim(x; method=LBFGS()) .^ 2 ≈ x #src + +# Forward mode autodiff + +ForwardDiff.jacobian(_x -> implicit_optim(_x; method=LBFGS()), x) +@test ForwardDiff.jacobian(_x -> implicit_optim(_x; method=LBFGS()), x) ≈ J #src + +#= +In this instance, we could use ForwardDiff.jl directly on the solver, but it returns the wrong result (not sure why). +=# + +ForwardDiff.jacobian(_x -> forward_optim(x; method=LBFGS()), x) + +# Reverse mode autodiff + +Zygote.jacobian(_x -> implicit_optim(_x; method=LBFGS()), x)[1] +@test Zygote.jacobian(_x -> implicit_optim(_x; method=LBFGS()), x)[1] ≈ J #src + +#= +In this instance, we cannot use Zygote.jl directly on the solver (due to unsupported `try/catch` statements). +=# + +try + Zygote.jacobian(_x -> forward_optim(x; method=LBFGS()), x)[1] +catch e + e +end + +# ## Nonlinear system + +#= +Next, we show how to differentiate through the solution of a nonlinear system of equations: +```math +\text{find} \quad y(x) \quad \text{such that} \quad F(x, y(x)) = 0 +``` +The optimality conditions are pretty obvious: +```math +F(x, y) = 0 +``` +=# + +#= +To make verification easy, we solve the following system: +```math +F(x, y) = y \odot y - x = 0 +``` +In this case, the optimization problem boils down to the componentwise square root function, but we implement it using a black box solver from [NLsolve.jl](https://github.com/JuliaNLSolvers/NLsolve.jl). +=# + +function forward_nlsolve(x; method) + F!(storage, y) = (storage .= y .^ 2 - x) + initial_y = similar(x) + initial_y .= 1 + result = nlsolve(F!, initial_y; method) + return result.zero +end + +#- + +function conditions_nlsolve(x, y; method) + c = y .^ 2 .- x + return c +end + +#- + +implicit_nlsolve = ImplicitFunction(forward_nlsolve, conditions_nlsolve) + +#- + +implicit_nlsolve(x; method=:newton) .^ 2 +@test implicit_nlsolve(x; method=:newton) .^ 2 ≈ x #src + +# Forward mode autodiff + +ForwardDiff.jacobian(_x -> implicit_nlsolve(_x; method=:newton), x) +@test ForwardDiff.jacobian(_x -> implicit_nlsolve(_x; method=:newton), x) ≈ J #src + +#- + +ForwardDiff.jacobian(_x -> forward_nlsolve(_x; method=:newton), x) + +# Reverse mode autodiff + +Zygote.jacobian(_x -> implicit_nlsolve(_x; method=:newton), x)[1] +@test Zygote.jacobian(_x -> implicit_nlsolve(_x; method=:newton), x)[1] ≈ J #src + +#- + +try + Zygote.jacobian(_x -> forward_nlsolve(_x; method=:newton), x)[1] +catch e + e +end + +# ## Fixed point + +#= +Finally, we show how to differentiate through the limit of a fixed point iteration: +```math +y \longmapsto T(x, y) +``` +The optimality conditions are pretty obvious: +```math +y = T(x, y) +``` +=# + +#= +To make verification easy, we consider [Heron's method](https://en.wikipedia.org/wiki/Methods_of_computing_square_roots#Heron's_method): +```math +T(x, y) = \frac{1}{2} \left(y + \frac{x}{y}\right) +``` +In this case, the fixed point algorithm boils down to the componentwise square root function, but we implement it manually. +=# + +function forward_fixedpoint(x; iterations) + y = ones(eltype(x), size(x)) + for _ in 1:iterations + y .= 0.5 .* (y .+ x ./ y) + end + return y +end + +#- + +function conditions_fixedpoint(x, y; iterations) + T = 0.5 .* (y .+ x ./ y) + return T .- y +end + +#- + +implicit_fixedpoint = ImplicitFunction(forward_fixedpoint, conditions_fixedpoint) + +#- + +implicit_fixedpoint(x; iterations=10) .^ 2 +@test implicit_fixedpoint(x; iterations=10) .^ 2 ≈ x #src + +# Forward mode autodiff + +ForwardDiff.jacobian(_x -> implicit_fixedpoint(_x; iterations=10), x) +@test ForwardDiff.jacobian(_x -> implicit_fixedpoint(_x; iterations=10), x) ≈ J #src + +#- + +ForwardDiff.jacobian(_x -> forward_fixedpoint(_x; iterations=10), x) + +# Reverse mode autodiff + +Zygote.jacobian(_x -> implicit_fixedpoint(_x; iterations=10), x)[1] +@test Zygote.jacobian(_x -> implicit_fixedpoint(_x; iterations=10), x)[1] ≈ J #src + +#- + +try + Zygote.jacobian(_x -> forward_fixedpoint(_x; iterations=10), x)[1] +catch e + e +end diff --git a/examples/1_unconstrained_optim.jl b/examples/1_unconstrained_optim.jl deleted file mode 100644 index 44b7707..0000000 --- a/examples/1_unconstrained_optim.jl +++ /dev/null @@ -1,99 +0,0 @@ -# # Unconstrained optimization - -#= -In this example, we show how to differentiate through the solution of an unconstrained optimization problem: -```math -y(x) = \underset{y \in \mathbb{R}^m}{\mathrm{argmin}} ~ f(x, y) -``` -The optimality conditions are given by gradient stationarity: -```math -\nabla_2 f(x, y) = 0 -``` -=# - -using ForwardDiff -using ImplicitDifferentiation -using LinearAlgebra -using Optim -using Random -using Test #src -using Zygote - -Random.seed!(63); - -# ## Implicit function - -#= -To make verification easy, we minimize the following objective: -```math -f(x, y) = \lVert y \odot y - x \rVert^2 -``` -In this case, the optimization problem boils down to the componentwise square root function, but we implement it using a black box solver from [Optim.jl](https://github.com/JuliaNLSolvers/Optim.jl). -Note the presence of a keyword argument. -=# - -function forward_optim(x; method) - f(y) = sum(abs2, y .^ 2 .- x) - y0 = ones(eltype(x), size(x)) - result = optimize(f, y0, method) - return Optim.minimizer(result) -end - -#= -Even though they are defined as a gradient, it is better to provide optimality conditions explicitly: that way we avoid nesting autodiff calls. By default, the conditions should accept two arguments as input. -The forward mapping and the conditions should accept the same set of keyword arguments. -=# - -function conditions_optim(x, y; method) - ∇₂f = 2 .* (y .^ 2 .- x) - return ∇₂f -end - -#= -We now have all the ingredients to construct our implicit function. -=# - -implicit_optim = ImplicitFunction(forward_optim, conditions_optim) - -# And indeed, it behaves as it should when we call it: - -x = rand(2) - -#- - -implicit_optim(x; method=LBFGS()) .^ 2 -@test implicit_optim(x; method=LBFGS()) .^ 2 ≈ x #src - -#= -Let's see what the explicit Jacobian looks like. -=# - -J = Diagonal(0.5 ./ sqrt.(x)) - -# ## Forward mode autodiff - -ForwardDiff.jacobian(_x -> implicit_optim(_x; method=LBFGS()), x) -@test ForwardDiff.jacobian(_x -> implicit_optim(_x; method=LBFGS()), x) ≈ J #src - -#= -Unsurprisingly, the Jacobian is the identity. -In this instance, we could use ForwardDiff.jl directly on the solver, but it returns the wrong result (not sure why). -=# - -ForwardDiff.jacobian(_x -> forward_optim(x; method=LBFGS()), x) - -# ## Reverse mode autodiff - -Zygote.jacobian(_x -> implicit_optim(_x; method=LBFGS()), x)[1] -@test Zygote.jacobian(_x -> implicit_optim(_x; method=LBFGS()), x)[1] ≈ J #src - -#= -Again, the Jacobian is the identity. -In this instance, we cannot use Zygote.jl directly on the solver (due to unsupported `try/catch` statements). -=# - -try - Zygote.jacobian(_x -> forward_optim(x; method=LBFGS()), x)[1] -catch e - e -end diff --git a/examples/4_constrained_optim.jl b/examples/2_advanced.jl similarity index 88% rename from examples/4_constrained_optim.jl rename to examples/2_advanced.jl index 6a59ab8..cc1e480 100644 --- a/examples/4_constrained_optim.jl +++ b/examples/2_advanced.jl @@ -1,15 +1,8 @@ -# # Constrained optimization +# # Advanced use cases #= -In this example, we show how to differentiate through the solution of a constrained optimization problem: -```math -y(x) = \underset{y \in \mathbb{R}^m}{\mathrm{argmin}} ~ f(x, y) \quad \text{subject to} \quad g(x, y) \leq 0 -``` -The optimality conditions are a bit trickier than in the previous cases. -We can projection on the feasible set $\mathcal{C}(x) = \{y: g(x, y) \leq 0 \}$ and exploit the convergence of projected gradient descent with step size $\eta$: -```math -y = \mathrm{proj}_{\mathcal{C}(x)} (y - \eta \nabla_2 f(x, y)) -``` +We dive into more advanced applications of implicit differentiation: +- constrained optimization problems =# using ForwardDiff @@ -22,7 +15,19 @@ using Zygote Random.seed!(63); -# ## Implicit function +# ## Constrained optimization + +#= +First, we show how to differentiate through the solution of a constrained optimization problem: +```math +y(x) = \underset{y \in \mathbb{R}^m}{\mathrm{argmin}} ~ f(x, y) \quad \text{subject to} \quad g(x, y) \leq 0 +``` +The optimality conditions are a bit trickier than in the previous cases. +We can projection on the feasible set $\mathcal{C}(x) = \{y: g(x, y) \leq 0 \}$ and exploit the convergence of projected gradient descent with step size $\eta$: +```math +y = \mathrm{proj}_{\mathcal{C}(x)} (y - \eta \nabla_2 f(x, y)) +``` +=# #= To make verification easy, we minimize the following objective: @@ -74,7 +79,7 @@ implicit_cstr_optim(x) .^ 2 J_thres = Diagonal([0.5 / sqrt(x[1]), 0]) -# ## Forward mode autodiff +# Forward mode autodiff ForwardDiff.jacobian(implicit_cstr_optim, x) @test ForwardDiff.jacobian(implicit_cstr_optim, x) ≈ J_thres #src @@ -83,7 +88,7 @@ ForwardDiff.jacobian(implicit_cstr_optim, x) ForwardDiff.jacobian(forward_cstr_optim, x) -# ## Reverse mode autodiff +# Reverse mode autodiff Zygote.jacobian(implicit_cstr_optim, x)[1] @test Zygote.jacobian(implicit_cstr_optim, x)[1] ≈ J_thres #src diff --git a/examples/2_nonlinear_solve.jl b/examples/2_nonlinear_solve.jl deleted file mode 100644 index b08f394..0000000 --- a/examples/2_nonlinear_solve.jl +++ /dev/null @@ -1,86 +0,0 @@ -# # Nonlinear solve - -#= -In this example, we show how to differentiate through the solution of a nonlinear system of equations: -```math -\text{find} \quad y(x) \quad \text{such that} \quad F(x, y(x)) = 0 -``` -The optimality conditions are pretty obvious: -```math -F(x, y) = 0 -``` -=# - -using ForwardDiff -using ImplicitDifferentiation -using LinearAlgebra -using NLsolve -using Random -using Test #src -using Zygote - -Random.seed!(63); - -# ## Implicit function - -#= -To make verification easy, we solve the following system: -```math -F(x, y) = y \odot y - x = 0 -``` -In this case, the optimization problem boils down to the componentwise square root function, but we implement it using a black box solver from [NLsolve.jl](https://github.com/JuliaNLSolvers/NLsolve.jl). -=# - -function forward_nlsolve(x; method) - F!(storage, y) = (storage .= y .^ 2 - x) - initial_y = similar(x) - initial_y .= 1 - result = nlsolve(F!, initial_y; method) - return result.zero -end - -#- - -function conditions_nlsolve(x, y; method) - c = y .^ 2 .- x - return c -end - -#- - -implicit_nlsolve = ImplicitFunction(forward_nlsolve, conditions_nlsolve) - -#- - -x = rand(2) - -#- - -implicit_nlsolve(x; method=:newton) .^ 2 -@test implicit_nlsolve(x; method=:newton) .^ 2 ≈ x #src - -#- - -J = Diagonal(0.5 ./ sqrt.(x)) - -# ## Forward mode autodiff - -ForwardDiff.jacobian(_x -> implicit_nlsolve(_x; method=:newton), x) -@test ForwardDiff.jacobian(_x -> implicit_nlsolve(_x; method=:newton), x) ≈ J #src - -#- - -ForwardDiff.jacobian(_x -> forward_nlsolve(_x; method=:newton), x) - -# ## Reverse mode autodiff - -Zygote.jacobian(_x -> implicit_nlsolve(_x; method=:newton), x)[1] -@test Zygote.jacobian(_x -> implicit_nlsolve(_x; method=:newton), x)[1] ≈ J #src - -#- - -try - Zygote.jacobian(_x -> forward_nlsolve(_x; method=:newton), x)[1] -catch e - e -end diff --git a/examples/3_fixed_points.jl b/examples/3_fixed_points.jl deleted file mode 100644 index 1741c5d..0000000 --- a/examples/3_fixed_points.jl +++ /dev/null @@ -1,85 +0,0 @@ -# # Fixed point - -#= -In this example, we show how to differentiate through the limit of a fixed point iteration: -```math -y \longmapsto T(x, y) -``` -The optimality conditions are pretty obvious: -```math -y = T(x, y) -``` -=# - -using ForwardDiff -using ImplicitDifferentiation -using LinearAlgebra -using Random -using Test #src -using Zygote - -Random.seed!(63); - -# ## Implicit function - -#= -To make verification easy, we consider [Heron's method](https://en.wikipedia.org/wiki/Methods_of_computing_square_roots#Heron's_method): -```math -T(x, y) = \frac{1}{2} \left(y + \frac{x}{y}\right) -``` -In this case, the fixed point algorithm boils down to the componentwise square root function, but we implement it manually. -=# - -function forward_fixedpoint(x; iterations) - y = ones(eltype(x), size(x)) - for _ in 1:iterations - y .= 0.5 .* (y .+ x ./ y) - end - return y -end - -#- - -function conditions_fixedpoint(x, y; iterations) - T = 0.5 .* (y .+ x ./ y) - return T .- y -end - -#- - -implicit_fixedpoint = ImplicitFunction(forward_fixedpoint, conditions_fixedpoint) - -#- - -x = rand(2) - -#- - -implicit_fixedpoint(x; iterations=10) .^ 2 -@test implicit_fixedpoint(x; iterations=10) .^ 2 ≈ x #src - -#- - -J = Diagonal(0.5 ./ sqrt.(x)) - -# ## Forward mode autodiff - -ForwardDiff.jacobian(_x -> implicit_fixedpoint(_x; iterations=10), x) -@test ForwardDiff.jacobian(_x -> implicit_fixedpoint(_x; iterations=10), x) ≈ J #src - -#- - -ForwardDiff.jacobian(_x -> forward_fixedpoint(_x; iterations=10), x) - -# ## Reverse mode autodiff - -Zygote.jacobian(_x -> implicit_fixedpoint(_x; iterations=10), x)[1] -@test Zygote.jacobian(_x -> implicit_fixedpoint(_x; iterations=10), x)[1] ≈ J #src - -#- - -try - Zygote.jacobian(_x -> forward_fixedpoint(_x; iterations=10), x)[1] -catch e - e -end diff --git a/examples/5_multiargs.jl b/examples/3_tricks.jl similarity index 51% rename from examples/5_multiargs.jl rename to examples/3_tricks.jl index 27c7301..ded836f 100644 --- a/examples/5_multiargs.jl +++ b/examples/3_tricks.jl @@ -1,11 +1,7 @@ -# # Multiple arguments +# # Tricks #= -In this example, we explain what to do when your function takes multiple input arguments: -```math -y(a, b) = \underset{y \in \mathbb{R}^m}{\mathrm{argmin}} ~ f(a, b, y) -``` -The key idea is to store both $a$ and $b$ inside a single vector $x$. +We demonstrate several features that may come in handy for some users. =# using ComponentArrays @@ -19,7 +15,15 @@ using Zygote Random.seed!(63); -# ## Implicit function +# ## Multiple arguments + +#= +First, we explain what to do when your forward mapping takes multiple input arguments: +```math +y(a, b) = \underset{y \in \mathbb{R}^m}{\mathrm{argmin}} ~ f(a, b, y) +``` +The key idea is to store both $a$ and $b$ inside a single vector $x$. +=# #= To make verification easy, we minimize the following objective: @@ -71,12 +75,76 @@ Let's see what the explicit Jacobian looks like. J = hcat(Diagonal(0.5 ./ sqrt.(x.a + 2x.b)), 2 * Diagonal(0.5 ./ sqrt.(x.a + 2x.b))) -# ## Forward mode autodiff +# Forward mode autodiff ForwardDiff.jacobian(implicit_components, x) @test ForwardDiff.jacobian(implicit_components, x) ≈ J #src -# ## Reverse mode autodiff +# Reverse mode autodiff Zygote.jacobian(implicit_components, x)[1] @test Zygote.jacobian(implicit_components, x)[1] ≈ J #src + +# ## Byproducts + +#= +Next, we explain what to do when your forward mapping computes another object that you want to keep track of, which we will call its "byproduct". +The difference between this and multiple outputs (which should be managed with ComponentArrays.jl) is that _we do not compute derivatives with respect to byproducts_. +=# + +#= +Imagine a situation where, depending on a coin toss, said mapping either doubles or halves the input. +After all, why not? +For each individual run, the algorithmic derivative is well-defined. +But to obtain it, you need to store the result of the toss. +=# + +function forward_cointoss(x) + z = rand(Bool) + if z + y = 2x + else + y = x / 2 + end + return y, z +end + +#= +And naturally, the optimality condition also depends on the toss. +=# + +function conditions_cointoss(x, y, z) + if z + return y .- 2x + else + return 2y .- x + end +end + +#= +To make sure that the implicit function you create takes this byproduct into account, just construct it like this: +=# + +implicit_cointoss = ImplicitFunction( + forward_cointoss, conditions_cointoss, HandleByproduct() +) + +#= +Then you have two ways of calling the function: the standard way will only return `y` +=# + +x = [1.0, 1.0] + +implicit_cointoss(x) + +#= +Or if you also need the byproduct, you can do +=# + +implicit_cointoss(x, ReturnByproduct()) + +#= +But whatever you choose, the byproduct is taken into account during differentiation! +=# + +Zygote.withjacobian(implicit_cointoss, x) From 689ff1c38638eb85d7954d0efbe75c83bc287ddf Mon Sep 17 00:00:00 2001 From: Guillaume Dalle <22795598+gdalle@users.noreply.github.com> Date: Fri, 28 Jul 2023 18:22:47 +0200 Subject: [PATCH 5/8] Reintroduce ignored docstrings --- docs/src/api.md | 5 +++++ src/implicit_function.jl | 26 ++++++++++---------------- src/utils.jl | 15 +++++++++++++++ 3 files changed, 30 insertions(+), 16 deletions(-) diff --git a/docs/src/api.md b/docs/src/api.md index dac41bc..2aea956 100644 --- a/docs/src/api.md +++ b/docs/src/api.md @@ -4,6 +4,10 @@ ```@docs ImplicitFunction +DirectLinearSolver +IterativeLinearSolver +HandleByproduct +ReturnByproduct ``` ## Internals @@ -11,6 +15,7 @@ ImplicitFunction ```@docs ImplicitDifferentiation.Forward ImplicitDifferentiation.Conditions +ImplicitDifferentiation.AbstractLinearSolver ImplicitDifferentiation.PushforwardMul! ImplicitDifferentiation.PullbackMul! ``` diff --git a/src/implicit_function.jl b/src/implicit_function.jl index 62d2a89..e2ac2de 100644 --- a/src/implicit_function.jl +++ b/src/implicit_function.jl @@ -5,14 +5,18 @@ Differentiable wrapper for an implicit function defined by a forward mapping and # Constructors - ImplicitFunction(f, c) - ImplicitFunction(f, c, linear_solver) - ImplicitFunction(f, c, HandleByproduct()) - ImplicitFunction(f, c, linear_solver, HandleByproduct()) - -Construct an `ImplicitFunction` from a forward mapping `f` and conditions `c`, both of which are Julia callables. +You can construct an `ImplicitFunction` from a forward mapping `f` and conditions `c`, both of which must be callables (function-like objects). While `f` does not not need to be compatible with automatic differentiation, `c` has to be. + ImplicitFunction(f, c[, HandleByproduct()]) + ImplicitFunction(f, c, linear_solver[, HandleByproduct()]) + +# Callable behavior + +An `ImplicitFunction` object `implicit` behaves like a function, and every call to it is differentiable. + + implicit(x::AbstractArray[, ReturnByproduct()]; kwargs...) + # Details - By default, we assume that the forward mapping is `x -> y(x)` and the conditions are `c(x,y(x)) = 0`. @@ -67,21 +71,11 @@ function Base.show(io::IO, implicit::ImplicitFunction) return print(io, "ImplicitFunction($(forward.f), $(conditions.c), $linear_solver)") end -""" - implicit(x::AbstractArray; kwargs...) - -Make an [`ImplicitFunction`](@ref) callable by applying the forward mapping `implicit.forward`. -""" function (implicit::ImplicitFunction)(x::AbstractArray; kwargs...) y, z = implicit.forward(x; kwargs...) return y end -""" - implicit(x::AbstractArray, ReturnByproduct(); kwargs...) - -Make an [`ImplicitFunction`](@ref) callable by applying the forward mapping `implicit.forward` and returning the byproduct. -""" function (implicit::ImplicitFunction)(x::AbstractArray, ::ReturnByproduct; kwargs...) y, z = implicit.forward(x, ; kwargs...) return (y, z) diff --git a/src/utils.jl b/src/utils.jl index 6ceb647..0cbfa87 100644 --- a/src/utils.jl +++ b/src/utils.jl @@ -1,4 +1,19 @@ +""" + HandleByproduct + +Trivial struct specifying that the forward mapping and conditions handle a byproduct. + +Used in the constructor for `ImplicitFunction`. +""" struct HandleByproduct end + +""" + ReturnByproduct + +Trivial struct specifying that we want to obtain a byproduct in addition to the solution. + +Used when calling an `ImplicitFunction`. +""" struct ReturnByproduct end """ From b18ea8cb8e97b7627d640679f3a261fd27d2db16 Mon Sep 17 00:00:00 2001 From: Guillaume Dalle <22795598+gdalle@users.noreply.github.com> Date: Sun, 30 Jul 2023 13:49:45 +0200 Subject: [PATCH 6/8] More tests for byproducts --- docs/Manifest.toml | 2 +- docs/Project.toml | 1 + docs/make.jl | 19 +- docs/src/api.md | 1 + docs/src/faq.md | 46 +++-- docs/src/index.md | 2 +- examples/1_basic.jl | 4 +- examples/2_advanced.jl | 2 +- ext/ImplicitDifferentiationChainRulesExt.jl | 6 +- ext/ImplicitDifferentiationForwardDiffExt.jl | 2 +- src/conditions.jl | 2 + src/forward.jl | 2 + src/implicit_function.jl | 6 +- src/utils.jl | 8 +- test/systematic.jl | 179 ++++++++++++++----- 15 files changed, 209 insertions(+), 73 deletions(-) diff --git a/docs/Manifest.toml b/docs/Manifest.toml index 9b9c6fd..0a2dbdb 100644 --- a/docs/Manifest.toml +++ b/docs/Manifest.toml @@ -2,7 +2,7 @@ julia_version = "1.9.2" manifest_format = "2.0" -project_hash = "e7f4896b7e8c3921c7466749f467ab7680867992" +project_hash = "ff84ddc3d5227f964f2cd507ce5cbc83b4fba207" [[deps.AMD]] deps = ["Libdl", "LinearAlgebra", "SparseArrays", "Test"] diff --git a/docs/Project.toml b/docs/Project.toml index aa54fc0..521326f 100644 --- a/docs/Project.toml +++ b/docs/Project.toml @@ -10,5 +10,6 @@ Literate = "98b081ad-f1c9-55d3-8b20-4c87d4299306" NLsolve = "2774e3e8-f4cf-5e23-947b-6d7e65073b56" Optim = "429524aa-4258-5aef-a3af-852621145aeb" Random = "9a3f8284-a2c9-5f02-9a11-845980a1fd5c" +StaticArrays = "90137ffa-7385-5640-81b9-e52037218182" Test = "8dfed614-e22c-5e08-85e1-65c5234f0b40" Zygote = "e88e6eb3-aa80-5325-afca-941959d7151f" diff --git a/docs/make.jl b/docs/make.jl index 515c1e1..5a8c714 100644 --- a/docs/make.jl +++ b/docs/make.jl @@ -3,6 +3,7 @@ using Documenter using ForwardDiff: ForwardDiff using ImplicitDifferentiation using Literate +using StaticArrays: StaticArrays DocMeta.setdocmeta!( ImplicitDifferentiation, :DocTestSetup, :(using ImplicitDifferentiation); recursive=true @@ -62,8 +63,24 @@ fmt = Documenter.HTML(; edit_link=:commit, ) +if isdefined(Base, :get_extension) + extension_modules = [ + Base.get_extension(ImplicitDifferentiation, :ImplicitDifferentiationChainRulesExt) + Base.get_extension(ImplicitDifferentiation, :ImplicitDifferentiationForwardDiffExt) + Base.get_extension( + ImplicitDifferentiation, :ImplicitDifferentiationStaticArraysExt + ) + ] +else + extension_modules = [ + ImplicitDifferentiation.ImplicitDifferentiationChainRulesExt, + ImplicitDifferentiation.ImplicitDifferentiationForwardDiffExt, + ImplicitDifferentiation.ImplicitDifferentiationStaticArraysExt, + ] +end + makedocs(; - modules=[ImplicitDifferentiation], + modules=vcat([ImplicitDifferentiation], extension_modules), authors="Guillaume Dalle, Mohamed Tarek and contributors", repo="https://github.com/gdalle/ImplicitDifferentiation.jl/blob/{commit}{path}#{line}", sitename="ImplicitDifferentiation.jl", diff --git a/docs/src/api.md b/docs/src/api.md index 2aea956..400d41a 100644 --- a/docs/src/api.md +++ b/docs/src/api.md @@ -8,6 +8,7 @@ DirectLinearSolver IterativeLinearSolver HandleByproduct ReturnByproduct +ChainRulesCore.rrule ``` ## Internals diff --git a/docs/src/faq.md b/docs/src/faq.md index d92d238..c9e4cf3 100644 --- a/docs/src/faq.md +++ b/docs/src/faq.md @@ -2,25 +2,36 @@ ## Supported autodiff backends -- Forward mode: [ForwardDiff.jl](https://github.com/JuliaDiff/ForwardDiff.jl) -- Reverse mode: all the packages compatible with [ChainRules.jl](https://github.com/JuliaDiff/ChainRules.jl) +| Mode | Backend | Support | +| ------- | ---------------------------------------------------------- | ------- | +| Forward | [ForwardDiff.jl] | yes | +| Reverse | [ChainRules.jl]-compatible ([Zygote.jl], [ReverseDiff.jl]) | yes | +| Forward | [ChainRules.jl]-compatible ([Diffractor.jl]) | soon | +| Both | [Enzyme.jl] | someday | -In the future, we would like to add +[ForwardDiff.jl]: https://github.com/JuliaDiff/ForwardDiff.jl +[ChainRules.jl]: https://github.com/JuliaDiff/ChainRules.jl +[Zygote.jl]: https://github.com/FluxML/Zygote.jl +[ReverseDiff.jl]: https://github.com/JuliaDiff/ReverseDiff.jl +[Enzyme.jl]: https://github.com/EnzymeAD/Enzyme.jl +[Diffractor.jl]: https://github.com/JuliaDiff/Diffractor.jl -- [Enzyme.jl](https://github.com/EnzymeAD/Enzyme.jl) -- [Diffractor.jl](https://github.com/JuliaDiff/Diffractor.jl) +## Writing conditions -## Higher-dimensional arrays +We recommend that the conditions themselves do not involve calls to autodiff, even when they describe a gradient. +Otherwise, you will need to make sure that nested autodiff works well in your case. +For instance, if you're differentiating your implicit function in reverse mode with Zygote.jl, you may want to use [`Zygote.forwarddiff`](https://fluxml.ai/Zygote.jl/stable/utils/#Zygote.forwarddiff) to wrap the conditions and differentiate them with ForwardDiff.jl instead. + +## Matrices and higher-order arrays For simplicity, our examples only display functions that eat and spit out vectors. However, arbitrary array shapes are supported, as long as the forward mapping _and_ conditions return similar arrays. Beware however, sparse arrays will be densified in the differentiation process. -## Scalar input / output +## Scalars Functions that eat or spit out a single number are not supported. -The forward mapping _and_ conditions need arrays: for example, instead of returning `value` you should return `[value]` (a 1-element `Vector`). -Consider using an `SVector` from [StaticArrays.jl](https://github.com/JuliaArrays/StaticArrays.jl) if you seek increased performance. +The forward mapping _and_ conditions need arrays: for example, instead of returning `val` you should return `[val]` (a 1-element `Vector`). ## Multiple inputs / outputs @@ -44,17 +55,18 @@ The same trick works for multiple outputs. ## Using byproducts -At first glance, it is not obvious why we impose that the forward mapping should return a byproduct `z` in addition to `y`. -It is mainly useful when the solution procedure creates objects such as Jacobians, which we want to reuse when computing or differentiating the `conditions`. -We will provide simple examples soon. -In the meantime, an advanced application is given by [DifferentiableFrankWolfe.jl](https://github.com/gdalle/DifferentiableFrankWolfe.jl). +Why would the forward mapping return a byproduct `z` in addition to `y`? +It is mainly useful when the solution procedure creates objects such as Jacobians, which we want to reuse when computing or differentiating the conditions. +In that case, you may want to write the differentiation rules yourself for the conditions. +A more advanced application is given by [DifferentiableFrankWolfe.jl](https://github.com/gdalle/DifferentiableFrankWolfe.jl). + +Keep in mind that derivatives of `z` will not be computed: the byproduct is considered constant during differentiation (unlike the case of multiple outputs outlined above). -## Differentiating byproducts +## Performance tips -Nope. Sorry. Don't even think about it. -The package is not designed to compute derivatives of `z`, only `y`, which is why the byproduct is considered constant during differentiation. +If you work with small arrays (say, less than 100 elements), consider using [StaticArrays.jl](https://github.com/JuliaArrays/StaticArrays.jl) if you seek increased performance. -## Modeling constrained optimization problems +## Modeling tips To express constrained optimization problems as implicit functions, you might need differentiable projections or proximal operators to write the optimality conditions. See [_Efficient and modular implicit differentiation_](https://arxiv.org/abs/2105.15183) for precise formulations. diff --git a/docs/src/index.md b/docs/src/index.md index be50c78..42c296d 100644 --- a/docs/src/index.md +++ b/docs/src/index.md @@ -7,7 +7,7 @@ CurrentModule = ImplicitDifferentiation [ImplicitDifferentiation.jl](https://github.com/gdalle/ImplicitDifferentiation.jl) is a package for automatic differentiation of functions defined implicitly, i.e., _forward mappings_ ```math -x \in \mathbb{R}^n \longmapsto y(x) \in \mathbb{R}^m +f: x \in \mathbb{R}^n \longmapsto y(x) \in \mathbb{R}^m ``` whose output is defined by _conditions_ diff --git a/examples/1_basic.jl b/examples/1_basic.jl index 69b5370..942e68f 100644 --- a/examples/1_basic.jl +++ b/examples/1_basic.jl @@ -66,7 +66,7 @@ The forward mapping and the conditions should accept the same set of keyword arg =# function conditions_optim(x, y; method) - ∇₂f = 2 .* (y .^ 2 .- x) + ∇₂f = @. 4 * (y^2 - x) * y return ∇₂f end @@ -129,7 +129,7 @@ In this case, the optimization problem boils down to the componentwise square ro =# function forward_nlsolve(x; method) - F!(storage, y) = (storage .= y .^ 2 - x) + F!(storage, y) = (storage .= y .^ 2 .- x) initial_y = similar(x) initial_y .= 1 result = nlsolve(F!, initial_y; method) diff --git a/examples/2_advanced.jl b/examples/2_advanced.jl index cc1e480..d1d48ff 100644 --- a/examples/2_advanced.jl +++ b/examples/2_advanced.jl @@ -55,7 +55,7 @@ function proj_hypercube(p) end function conditions_cstr_optim(x, y) - ∇₂f = 2 .* (y .^ 2 .- x) + ∇₂f = @. 4 * (y^2 - x) * y η = 0.1 return y .- proj_hypercube(y .- η .* ∇₂f) end diff --git a/ext/ImplicitDifferentiationChainRulesExt.jl b/ext/ImplicitDifferentiationChainRulesExt.jl index 4ee48b4..cfedb03 100644 --- a/ext/ImplicitDifferentiationChainRulesExt.jl +++ b/ext/ImplicitDifferentiationChainRulesExt.jl @@ -19,7 +19,7 @@ This is only available if ChainRulesCore.jl is loaded (extension), except on Jul - By default, this returns a single output `y(x)` with a pullback accepting a single cotangent `dy`. - If `ReturnByproduct()` is passed as an argument, this returns a couple of outputs `(y(x),z(x))` with a pullback accepting a couple of cotangents `(dy, dz)` (remember that `z(x)` is not differentiated so its cotangent is ignored). -We compute the vector-Jacobian product `Jᵀv` by solving `Aᵀu = v` and setting `Jᵀv = -Bᵀu`. +We compute the vector-Jacobian product `Jᵀv` by solving `Aᵀu = v` and setting `Jᵀv = -Bᵀu` (see [`ImplicitFunction`](@ref) for the definition of `A` and `B`). Keyword arguments are given to both `implicit.forward` and `implicit.conditions`. """ function ChainRulesCore.rrule( @@ -53,7 +53,7 @@ function ChainRulesCore.rrule( rc::RuleConfig, implicit::ImplicitFunction, x::AbstractArray{R}; kwargs... ) where {R} (y, z), implicit_pullback = rrule(rc, implicit, x, ReturnByproduct(); kwargs...) - implicit_pullback_no_byproduct(dy) = implicit_pullback((dy, nothing)) + implicit_pullback_no_byproduct(dy) = Base.front(implicit_pullback((dy, nothing))) return y, implicit_pullback_no_byproduct end @@ -73,7 +73,7 @@ function (implicit_pullback::ImplicitPullback)((dy, dz)) mul!(dx_vec, Bᵀ_op, dF_vec) lmul!(-one(R), dx_vec) dx = reshape(dx_vec, size(x)) - return (NoTangent(), dx) + return (NoTangent(), dx, NoTangent()) end end diff --git a/ext/ImplicitDifferentiationForwardDiffExt.jl b/ext/ImplicitDifferentiationForwardDiffExt.jl index 2bf89b0..b6bd98e 100644 --- a/ext/ImplicitDifferentiationForwardDiffExt.jl +++ b/ext/ImplicitDifferentiationForwardDiffExt.jl @@ -25,7 +25,7 @@ This is only available if ForwardDiff.jl is loaded (extension). - By default, this returns a single output `y_and_dy(x)`. - If `ReturnByproduct()` is passed as an argument, this returns a couple of outputs `(y_and_dy(x),z(x))` (remember that `z(x)` is not differentiated so `dz(x)` doesn't exist). -We compute the Jacobian-vector product `Jv` by solving `Au = Bv` and setting `Jv = u`. +We compute the Jacobian-vector product `Jv` by solving `Au = Bv` and setting `Jv = u` (see [`ImplicitFunction`](@ref) for the definition of `A` and `B`). Keyword arguments are given to both `implicit.forward` and `implicit.conditions`. """ function (implicit::ImplicitFunction)( diff --git a/src/conditions.jl b/src/conditions.jl index d0960bc..1d3d725 100644 --- a/src/conditions.jl +++ b/src/conditions.jl @@ -18,3 +18,5 @@ end (conditions::Conditions{true})(x, y, z; kwargs...) = conditions.c(x, y, z; kwargs...) (conditions::Conditions{false})(x, y, z; kwargs...) = conditions.c(x, y; kwargs...) + +handles_byproduct(::Conditions{byproduct}) where {byproduct} = byproduct diff --git a/src/forward.jl b/src/forward.jl index 06bb62c..97fa2ba 100644 --- a/src/forward.jl +++ b/src/forward.jl @@ -26,3 +26,5 @@ function (forward::Forward{false})(x; kwargs...) z = 0 return y, z end + +handles_byproduct(::Forward{byproduct}) where {byproduct} = byproduct diff --git a/src/implicit_function.jl b/src/implicit_function.jl index e2ac2de..2bf2465 100644 --- a/src/implicit_function.jl +++ b/src/implicit_function.jl @@ -77,6 +77,10 @@ function (implicit::ImplicitFunction)(x::AbstractArray; kwargs...) end function (implicit::ImplicitFunction)(x::AbstractArray, ::ReturnByproduct; kwargs...) - y, z = implicit.forward(x, ; kwargs...) + y, z = implicit.forward(x; kwargs...) return (y, z) end + +function handles_byproduct(implicit::ImplicitFunction) + return handles_byproduct(implicit.forward) && handles_byproduct(implicit.conditions) +end diff --git a/src/utils.jl b/src/utils.jl index 0cbfa87..3ef77ca 100644 --- a/src/utils.jl +++ b/src/utils.jl @@ -47,13 +47,17 @@ end function (pfm::PushforwardMul!)(res::AbstractVector, δinput_vec::AbstractVector) δinput = reshape(δinput_vec, pfm.input_size) δoutput = only(pfm.pushforward(δinput)) - return res .= vec(δoutput) + for i in eachindex(IndexLinear(), res, δoutput) + res[i] = δoutput[i] + end end function (pbm::PullbackMul!)(res::AbstractVector, δoutput_vec::AbstractVector) δoutput = reshape(δoutput_vec, pbm.output_size) δinput = only(pbm.pullback(δoutput)) - return res .= vec(δinput) + for i in eachindex(IndexLinear(), res, δinput) + res[i] = δinput[i] + end end ## Override this function from LinearOperators to avoid generating the whole methods table diff --git a/test/systematic.jl b/test/systematic.jl index bb5fce1..f6b2df7 100644 --- a/test/systematic.jl +++ b/test/systematic.jl @@ -2,26 +2,42 @@ using ChainRulesCore using ChainRulesTestUtils using ForwardDiff using ImplicitDifferentiation +using ImplicitDifferentiation: handles_byproduct using JET using LinearAlgebra using Random using StaticArrays using Test using Zygote +using Zygote: ZygoteRuleConfig Random.seed!(63); -""" - mysqrt(x) +function is_static_array(a) + return ( + typeof(a) <: StaticArray || + typeof(a) <: (Base.ReshapedArray{T,N,<:StaticArray} where {T,N}) + ) +end -Compute the elementwise square root, breaking Zygote.jl and ForwardDiff.jl in the process. -""" -function mysqrt(x::AbstractArray) +function break_forwarddiff_zygote(x) a = [0.0] - a[1] = first(x) + a[1] = float(first(x)) + return nothing +end + +function mysqrt(x::AbstractArray) + break_forwarddiff_zygote(x) return sqrt.(x) end +function mysqrt_byproduct(x::AbstractArray) + break_forwarddiff_zygote(x) + z = rand((2,)) + y = x .^ (1 / z) + return y, z +end + function make_implicit_sqrt(linear_solver) forward(x) = mysqrt(x) conditions(x, y) = y .^ 2 .- x @@ -29,63 +45,140 @@ function make_implicit_sqrt(linear_solver) return implicit end -function test_implicit_sqrt_call(implicit, x) - @test (@inferred implicit(x)) ≈ sqrt.(x) +function make_implicit_sqrt_byproduct(linear_solver) + forward(x) = mysqrt_byproduct(x) + conditions(x, y, z) = y .^ z .- x + implicit = ImplicitFunction(forward, conditions, linear_solver, HandleByproduct()) + return implicit +end + +function test_implicit_call(implicit, x; y_true) + @test_throws MethodError implicit("hello") + @test_throws MethodError implicit(x, x) + y1 = @inferred implicit(x) + y2, z2 = @inferred implicit(x, ReturnByproduct()) + @test y1 ≈ y_true + @test y2 ≈ y_true if typeof(x) <: StaticArray - @test typeof(implicit(x)) <: StaticArray + @test is_static_array(y1) + @test is_static_array(y2) + end + if handles_byproduct(implicit) + @test z2 == 2 + else + @test z2 == 0 end if VERSION >= v"1.9" - test_opt(implicit, (typeof(x),)) + @test_opt target_modules = (ImplicitDifferentiation,) implicit(x) + @test_call target_modules = (ImplicitDifferentiation,) implicit(x) end end -function test_implicit_sqrt_forward(implicit, x) - y = implicit(x) - J = Diagonal(0.5 ./ vec(sqrt.(x))) - @test ForwardDiff.jacobian(implicit, x) ≈ J +function test_implicit_forward(implicit, x; y_true, J_true) + # High-level + J1 = ForwardDiff.jacobian(implicit, x) + J2 = ForwardDiff.jacobian(x -> implicit(x, ReturnByproduct())[1], x) + @test J1 ≈ J_true + @test J2 ≈ J_true + # Low-level x_and_dx = ForwardDiff.Dual.(x, ((0, 0),)) - y_and_dy = @inferred implicit(x_and_dx) - @test size(y_and_dy) == size(y) + y_and_dy1 = @inferred implicit(x_and_dx) + y_and_dy2, z2 = @inferred implicit(x_and_dx, ReturnByproduct()) + @test size(y_and_dy1) == size(y_true) + @test size(y_and_dy2) == size(y_true) + @test ForwardDiff.value.(y_and_dy1) ≈ y_true + @test ForwardDiff.value.(y_and_dy2) ≈ y_true if typeof(x) <: StaticArray - @test typeof(y_and_dy) <: StaticArray || - typeof(y_and_dy) <: (Base.ReshapedArray{T,N,<:StaticArray} where {T,N}) + @test is_static_array(y_and_dy1) + @test is_static_array(y_and_dy2) + end + if handles_byproduct(implicit) + @test z2 == 2 + else + @test z2 == 0 + end + if VERSION >= v"1.9" + @test_skip @test_opt target_modules = (ImplicitDifferentiation,) implicit(x_and_dx) + @test_call target_modules = (ImplicitDifferentiation,) implicit(x_and_dx) end end -function test_implicit_sqrt_reverse(implicit, x) - J = Diagonal(0.5 ./ vec(sqrt.(x))) - @test Zygote.jacobian(implicit, x)[1] ≈ J - _, pullback = @inferred rrule(Zygote.ZygoteRuleConfig(), implicit, x) - dy = similar(implicit(x)) - dy .= zero(eltype(dy)) - @test (@inferred pullback(dy)) == pullback(dy) - _, dx = pullback(dy) - @test size(dx) == size(x) +function test_implicit_reverse(implicit, x; y_true, J_true) + # High-level + J1 = Zygote.jacobian(implicit, x)[1] + J2 = Zygote.jacobian(x -> implicit(x, ReturnByproduct())[1], x)[1] + @test J1 ≈ J_true + @test J2 ≈ J_true + # Low-level + y1, pb1 = @inferred rrule(ZygoteRuleConfig(), implicit, x) + (y2, z2), pb2 = @inferred rrule(ZygoteRuleConfig(), implicit, x, ReturnByproduct()) + @test y1 ≈ y_true + @test y2 ≈ y_true + dy1 = zeros(eltype(y1), size(y1)...) + dy2 = zeros(eltype(y2), size(y2)...) + dz2 = nothing + dimp1, dx1 = @inferred pb1(dy1) + dimp2, dx2, drp = @inferred pb2((dy2, dz2)) + @test size(dx1) == size(x) + @test size(dx2) == size(x) if typeof(x) <: StaticArray - @test typeof(dx) <: StaticArray || - typeof(dx) <: (Base.ReshapedArray{T,N,<:StaticArray} where {T,N}) + @test is_static_array(y1) + @test is_static_array(y2) + @test is_static_array(dx1) + @test is_static_array(dx2) + end + @test dimp1 isa NoTangent + @test dimp2 isa NoTangent + @test drp isa NoTangent + if handles_byproduct(implicit) + @test z2 == 2 + else + @test z2 == 0 + end + if VERSION >= v"1.9" + @test_skip @test_opt target_modules = (ImplicitDifferentiation,) rrule( + ZygoteRuleConfig(), implicit, x + ) + @test_skip @test_opt target_modules = (ImplicitDifferentiation,) pb1(dy1) + @test_call target_modules = (ImplicitDifferentiation,) rrule( + ZygoteRuleConfig(), implicit, x + ) + @test_call target_modules = (ImplicitDifferentiation,) pb1(dy1) end # Skipped because of https://github.com/JuliaDiff/ChainRulesTestUtils.jl/issues/232 and because it detects weird type instabilities @test_skip test_rrule(implicit, x) + @test_skip test_rrule(x -> implicit(x, ReturnByproduct()), x) end x_candidates = ( rand(2), rand(2, 3, 4), SVector{2}(rand(2)), SArray{Tuple{2,3,4}}(rand(2, 3, 4)) -) +); -for linear_solver in (IterativeLinearSolver(), DirectLinearSolver()) - forward(x) = mysqrt(x) - conditions(x, y) = y .^ 2 .- x - implicit = ImplicitFunction(forward, conditions, linear_solver) - for x in x_candidates - testsetname = "$(typeof(x)) - $(typeof(linear_solver))" - @testset verbose = true "$testsetname" begin - if x isa StaticArray && linear_solver isa IterativeLinearSolver - continue - end - test_implicit_sqrt_call(implicit, x) - test_implicit_sqrt_forward(implicit, x) - test_implicit_sqrt_reverse(implicit, x) +linear_solver_candidates = (IterativeLinearSolver(), DirectLinearSolver()) + +for linear_solver in linear_solver_candidates, x in x_candidates + if x isa StaticArray && linear_solver isa IterativeLinearSolver + continue + end + y_true = sqrt.(x) + J_true = Diagonal(0.5 ./ vec(sqrt.(x))) + + testsetname = "$(typeof(x)) - $(typeof(linear_solver))" + implicit_sqrt = make_implicit_sqrt(linear_solver) + implicit_sqrt_byproduct = make_implicit_sqrt_byproduct(linear_solver) + + @testset verbose = true "$testsetname" begin + @testset "Call" begin + test_implicit_call(implicit_sqrt, x; y_true) + test_implicit_call(implicit_sqrt_byproduct, x; y_true) + end + @testset "Forward" begin + test_implicit_forward(implicit_sqrt, x; y_true, J_true) + test_implicit_forward(implicit_sqrt_byproduct, x; y_true, J_true) + end + @testset "Reverse" begin + test_implicit_reverse(implicit_sqrt, x; y_true, J_true) + test_implicit_reverse(implicit_sqrt_byproduct, x; y_true, J_true) end end end From d38ec88a108517a75adb7d44750cec8813fc1de3 Mon Sep 17 00:00:00 2001 From: Mohamed Tarek Date: Sun, 30 Jul 2023 22:37:29 +1000 Subject: [PATCH 7/8] fix type instability --- ext/ImplicitDifferentiationForwardDiffExt.jl | 2 +- test/systematic.jl | 2 +- 2 files changed, 2 insertions(+), 2 deletions(-) diff --git a/ext/ImplicitDifferentiationForwardDiffExt.jl b/ext/ImplicitDifferentiationForwardDiffExt.jl index b6bd98e..ae62b28 100644 --- a/ext/ImplicitDifferentiationForwardDiffExt.jl +++ b/ext/ImplicitDifferentiationForwardDiffExt.jl @@ -45,7 +45,7 @@ function (implicit::ImplicitFunction)( B_op = LinearOperator(R, m, n, false, false, PushforwardMul!(pfB, size(x))) A_op_presolved = presolve(linear_solver, A_op, y) - dy = map(1:N) do k + dy = ntuple(Val(N)) do k dₖx_vec = vec(partials.(x_and_dx, k)) Bdx = vec(similar(y)) mul!(Bdx, B_op, dₖx_vec) diff --git a/test/systematic.jl b/test/systematic.jl index f6b2df7..4c103c3 100644 --- a/test/systematic.jl +++ b/test/systematic.jl @@ -98,7 +98,7 @@ function test_implicit_forward(implicit, x; y_true, J_true) @test z2 == 0 end if VERSION >= v"1.9" - @test_skip @test_opt target_modules = (ImplicitDifferentiation,) implicit(x_and_dx) + @test_opt target_modules = (ImplicitDifferentiation,) implicit(x_and_dx) @test_call target_modules = (ImplicitDifferentiation,) implicit(x_and_dx) end end From c6a0459d21bb54b84f9f261f6eea2cdb3b39cd03 Mon Sep 17 00:00:00 2001 From: Mohamed Tarek Date: Sun, 30 Jul 2023 22:54:59 +1000 Subject: [PATCH 8/8] fix CI --- test/systematic.jl | 39 +++++++++++++++++++++------------------ 1 file changed, 21 insertions(+), 18 deletions(-) diff --git a/test/systematic.jl b/test/systematic.jl index 4c103c3..568026d 100644 --- a/test/systematic.jl +++ b/test/systematic.jl @@ -11,6 +11,15 @@ using Test using Zygote using Zygote: ZygoteRuleConfig +@static if VERSION < v"1.9" + macro test_opt(x...) + return :() + end + macro test_call(x...) + return :() + end +end + Random.seed!(63); function is_static_array(a) @@ -68,10 +77,8 @@ function test_implicit_call(implicit, x; y_true) else @test z2 == 0 end - if VERSION >= v"1.9" - @test_opt target_modules = (ImplicitDifferentiation,) implicit(x) - @test_call target_modules = (ImplicitDifferentiation,) implicit(x) - end + @test_opt target_modules = (ImplicitDifferentiation,) implicit(x) + @test_call target_modules = (ImplicitDifferentiation,) implicit(x) end function test_implicit_forward(implicit, x; y_true, J_true) @@ -97,10 +104,8 @@ function test_implicit_forward(implicit, x; y_true, J_true) else @test z2 == 0 end - if VERSION >= v"1.9" - @test_opt target_modules = (ImplicitDifferentiation,) implicit(x_and_dx) - @test_call target_modules = (ImplicitDifferentiation,) implicit(x_and_dx) - end + @test_opt target_modules = (ImplicitDifferentiation,) implicit(x_and_dx) + @test_call target_modules = (ImplicitDifferentiation,) implicit(x_and_dx) end function test_implicit_reverse(implicit, x; y_true, J_true) @@ -135,16 +140,14 @@ function test_implicit_reverse(implicit, x; y_true, J_true) else @test z2 == 0 end - if VERSION >= v"1.9" - @test_skip @test_opt target_modules = (ImplicitDifferentiation,) rrule( - ZygoteRuleConfig(), implicit, x - ) - @test_skip @test_opt target_modules = (ImplicitDifferentiation,) pb1(dy1) - @test_call target_modules = (ImplicitDifferentiation,) rrule( - ZygoteRuleConfig(), implicit, x - ) - @test_call target_modules = (ImplicitDifferentiation,) pb1(dy1) - end + @test_skip @test_opt target_modules = (ImplicitDifferentiation,) rrule( + ZygoteRuleConfig(), implicit, x + ) + @test_skip @test_opt target_modules = (ImplicitDifferentiation,) pb1(dy1) + @test_call target_modules = (ImplicitDifferentiation,) rrule( + ZygoteRuleConfig(), implicit, x + ) + @test_call target_modules = (ImplicitDifferentiation,) pb1(dy1) # Skipped because of https://github.com/JuliaDiff/ChainRulesTestUtils.jl/issues/232 and because it detects weird type instabilities @test_skip test_rrule(implicit, x) @test_skip test_rrule(x -> implicit(x, ReturnByproduct()), x)