Simplify interface and add more tests for first release

Author: gdalle

Project.toml

@@ -5,10 +5,12 @@ version = "0.1.0"
 
 [deps]
 ChainRulesCore = "d360d2e6-b24c-11e9-a2a3-2a2ae2dbcce4"
+Krylov = "ba0b0d4f-ebba-5204-a429-3ac8c609bfb7"
 LinearOperators = "5c8ed15e-5a4c-59e4-a42b-c7e8811fb125"
 SparseArrays = "2f01184e-e22b-5df5-ae63-d93ebab69eaf"
 
 [compat]
 ChainRulesCore = "1.14"
+Krylov = "0.8.1"
 LinearOperators = "2.2.3"
 julia = "1.7"

README.md

@@ -7,6 +7,7 @@
 [![Code Style: Blue](https://img.shields.io/badge/code%20style-blue-4495d1.svg)](https://github.com/invenia/BlueStyle)
 [![Aqua QA](https://raw.githubusercontent.com/JuliaTesting/Aqua.jl/master/badge.svg)](https://github.com/JuliaTesting/Aqua.jl)
 
-Automatic differentiation of implicit functions. See the [documentation](https://gdalle.github.io/ImplicitDifferentiation.jl/dev) for details.
+Automatic differentiation of implicit functions.
+See the [documentation](https://gdalle.github.io/ImplicitDifferentiation.jl/dev) for details.
 
 > This package is in a very early development stage, so use it at your own risk!

docs/src/index.md

@@ -11,4 +11,8 @@ CurrentModule = ImplicitDifferentiation
 To install it, open a Julia Pkg REPL and run:
 ```julia
 pkg> add "https://github.com/gdalle/ImplicitDifferentiation.jl"
-```
+```
+
+## Related packages
+
+- [DiffOpt.jl](https://github.com/jump-dev/DiffOpt.jl): differentiation of convex optimization problems

src/ImplicitDifferentiation.jl

@@ -1,6 +1,7 @@
 module ImplicitDifferentiation
 
 using ChainRulesCore
+using Krylov
 using LinearOperators
 using SparseArrays
 

src/implicit_function.jl

@@ -14,12 +14,24 @@ If `x ∈ ℝⁿ`, `y ∈ ℝᵐ` and `F(x,y) ∈ ℝᶜ`, this amounts to solvi
 - `conditions::C`: callable of the form `(x,y) -> F(x,y)`
 - `linear_solver::L`: callable of the form `(A,b) -> u` such that `A * u = b`
 """
-Base.@kwdef struct ImplicitFunction{F,C,L}
+struct ImplicitFunction{F,C,L}
     forward::F
     conditions::C
     linear_solver::L
 end
 
+"""
+    ImplicitFunction(forward, conditions)
+
+Construct an `ImplicitFunction` with `Krylov.gmres` as the default linear solver.
+
+# See also
+- [`ImplicitFunction{F,C,L}`](@ref)
+"""
+function ImplicitFunction(forward::F, conditions::C) where {F,C}
+    return ImplicitFunction(forward, conditions, Krylov.gmres)
+end
+
 struct SolverFailureException <: Exception
     msg::String
 end
@@ -29,7 +41,7 @@ end
 
 Make [`ImplicitFunction`](@ref) callable by applying `implicit.forward`.
 """
-(implicit::ImplicitFunction)(x) = first(implicit.forward(x))
+(implicit::ImplicitFunction)(x) = implicit.forward(x)
 
 """
     frule(rc, (_, dx), implicit, x)
@@ -43,11 +55,11 @@ function ChainRulesCore.frule(
 )
     (; forward, conditions, linear_solver) = implicit
 
-    y, useful_info = forward(x)
+    y = forward(x)
     n, m = length(x), length(y)
 
-    conditions_x(x̃) = conditions(x̃, y, useful_info)
-    conditions_y(ỹ) = -conditions(x, ỹ, useful_info)
+    conditions_x(x̃) = conditions(x̃, y)
+    conditions_y(ỹ) = -conditions(x, ỹ)
 
     pushforward_A(dỹ) = frule_via_ad(rc, (NoTangent(), dỹ), conditions_y, y)[2]
     pushforward_B(dx̃) = frule_via_ad(rc, (NoTangent(), dx̃), conditions_x, x)[2]
@@ -77,11 +89,11 @@ We compute the vector-Jacobian product `Jᵀv` by solving `Aᵀu = v` and settin
 function ChainRulesCore.rrule(rc::RuleConfig, implicit::ImplicitFunction, x::AbstractVector)
     (; forward, conditions, linear_solver) = implicit
 
-    y, useful_info = forward(x)
+    y = forward(x)
     n, m = length(x), length(y)
 
-    conditions_x(x̃) = conditions(x̃, y, useful_info)
-    conditions_y(ỹ) = -conditions(x, ỹ, useful_info)
+    conditions_x(x̃) = conditions(x̃, y)
+    conditions_y(ỹ) = -conditions(x, ỹ)
 
     pullback_Aᵀ = last ∘ rrule_via_ad(rc, conditions_y, y)[2]
     pullback_Bᵀ = last ∘ rrule_via_ad(rc, conditions_x, x)[2]

src/simplex.jl

@@ -5,14 +5,18 @@ Compute the Euclidean projection onto the probability simplex and the set of ind
 
 See <https://arxiv.org/abs/1602.02068> for details.
 """
-function simplex_projection_and_support(z::AbstractVector{<:Real})
+function simplex_projection_and_support(z::AbstractVector{R}) where {R<:Real}
     d = length(z)
     z_sorted = sort(z; rev=true)
     z_sorted_cumsum = cumsum(z_sorted)
     k = maximum(j for j in 1:d if (1 + j * z_sorted[j]) > z_sorted_cumsum[j])
     τ = (z_sorted_cumsum[k] - 1) / k
-    p = max.(z .- τ, 0)
-    s = [Int(p[i] > eps()) for i in 1:d]
+    p = Vector{R}(undef, d)
+    s = Vector{Int}(undef, d)
+    for i in 1:d
+        p[i] = max(z[i] - τ, zero(R))
+        s[i] = Int(!iszero(p[i]))
+    end
     return p, s
 end;
 

test/1_unconstrained_optimization.jl

@@ -13,7 +13,6 @@ F(x, \hat{y}(x)) = 0 \quad \text{with} \quad F(x,y) = \nabla_2 f(x, y) = 0
 =#
 
 using ImplicitDifferentiation
-using Krylov: gmres
 using Optim: optimize, minimizer, LBFGS
 using Zygote
 
@@ -35,18 +34,18 @@ function forward(x)
     y0 = zero(x)
     res = optimize(f, y0, LBFGS(); autodiff=:forward)
     y = minimizer(res)
-    return y, nothing
+    return y
 end;
 
 #=
 On the other hand, optimality conditions should be provided explicitly whenever possible, so as to avoid nesting automatic differentiation calls.
 =#
 
-conditions(x, y, useful_info=nothing) = 2(y - x);
+conditions(x, y) = 2(y - x);
 
 # We now have all the ingredients to construct our implicit function.
 
-implicit = ImplicitFunction(; forward=forward, conditions=conditions, linear_solver=gmres);
+implicit = ImplicitFunction(forward, conditions);
 
 # ## Testing
 

test/2_constrained_optimization.jl

@@ -15,7 +15,6 @@ F(x, \hat{y}(x)) = 0 \quad \text{with} \quad F(x,y) = \mathrm{proj}_{\mathcal{C}
 using ImplicitDifferentiation
 using Ipopt
 using JuMP
-using Krylov: gmres
 using Zygote
 
 using ChainRulesTestUtils  #src
@@ -44,12 +43,12 @@ function forward(x)
     @constraint(model, sum(y) == 1)
     @objective(model, Min, sum((y .- x) .^ 2))
     optimize!(model)
-    return value.(y), nothing
+    return value.(y)
 end;
 
-conditions(x, y, useful_info=nothing) = simplex_projection(y - 0.1 * 2(y - x)) - y;
+conditions(x, y) = simplex_projection(y - 0.1 * 2(y - x)) - y;
 
-implicit = ImplicitFunction(; forward=forward, conditions=conditions, linear_solver=gmres);
+implicit = ImplicitFunction(forward, conditions);
 
 # ## Testing
 

test/3_optimal_transport.jl

@@ -4,12 +4,11 @@
 In this example, we show how to differentiate through the solution of the entropy-regularized optimal transport problem.
 =#
 
+using Distances
+using FiniteDiff
 using ImplicitDifferentiation
 using OptimalTransport
-using Distances
-using Krylov: gmres
 using Zygote
-using FiniteDiff
 using Test, LinearAlgebra #src
 
 #=
@@ -84,17 +83,17 @@ function forward(C_vec)
     solver = OptimalTransport.build_solver(μ, ν, C, ε, SinkhornGibbs())
     OptimalTransport.solve!(solver)
     û = solver.cache.u
-    return û, nothing
+    return û
 end
 
-function conditions(C_vec, û, useful_info=nothing)
+function conditions(C_vec, û)
     C = reshape(C_vec, n, m)
     K = exp.(.-C ./ ε)
     v̂ = ν ./ (K' * û)
     return û .- μ ./ (K * v̂)
 end
 
-implicit = ImplicitFunction(; forward=forward, conditions=conditions, linear_solver=gmres);
+implicit = ImplicitFunction(forward, conditions);
 
 # ## Testing
 

test/4_struct.jl

@@ -6,7 +6,6 @@ In this example, we demonstrate implicit differentiation through functions that
 
 using ComponentArrays
 using ImplicitDifferentiation
-using Krylov: gmres
 using Zygote
 
 using ChainRulesCore  #src
@@ -19,13 +18,13 @@ using Test  #src
 We replicate a componentwise square function with `NamedTuple`s, taking `a=(x,y)` as input and returning `b=(u,v)`.
 =#
 
-forward(a::ComponentVector) = ComponentVector(u=a.x .^ 2, v=a.y .^ 2), nothing;
+forward(a::ComponentVector) = ComponentVector(u=a.x .^ 2, v=a.y .^ 2)
 
-function conditions(a::ComponentVector, b::ComponentVector, useful_info=nothing)
+function conditions(a::ComponentVector, b::ComponentVector)
     return vcat(b.u .- a.x .^ 2, b.v .- a.y .^ 2)
 end
 
-implicit = ImplicitFunction(; forward=forward, conditions=conditions, linear_solver=gmres);
+implicit = ImplicitFunction(forward, conditions);
 
 #=
 In order to be able to call `Zygote.gradient`, we use `implicit` to define a convoluted version of the squared Euclidean norm, which takes a `ComponentVector` as input and returns a real number.

test/Manifest.toml

@@ -1,8 +1,14 @@
 # This file is machine-generated - editing it directly is not advised
 
-julia_version = "1.7.2"
+julia_version = "1.7.3"
 manifest_format = "2.0"
 
+[[deps.AMD]]
+deps = ["Libdl", "LinearAlgebra", "SparseArrays", "Test"]
+git-tree-sha1 = "fc66ffc5cff568936649445f58a55b81eaf9592c"
+uuid = "14f7f29c-3bd6-536c-9a0b-7339e30b5a3e"
+version = "0.4.0"
+
 [[deps.ASL_jll]]
 deps = ["Artifacts", "JLLWrappers", "Libdl", "Pkg"]
 git-tree-sha1 = "6252039f98492252f9e47c312c8ffda0e3b9e78d"
@@ -90,6 +96,12 @@ git-tree-sha1 = "bf98fa45a0a4cee295de98d4c1462be26345b9a1"
 uuid = "9e997f8a-9a97-42d5-a9f1-ce6bfc15e2c0"
 version = "0.1.2"
 
+[[deps.CodeTracking]]
+deps = ["InteractiveUtils", "UUIDs"]
+git-tree-sha1 = "6d4fa04343a7fc9f9cb9cff9558929f3d2752717"
+uuid = "da1fd8a2-8d9e-5ec2-8556-3022fb5608a2"
+version = "1.0.9"
+
 [[deps.CodecBzip2]]
 deps = ["Bzip2_jll", "Libdl", "TranscodingStreams"]
 git-tree-sha1 = "2e62a725210ce3c3c2e1a3080190e7ca491f18d7"
@@ -184,9 +196,22 @@ uuid = "ffbed154-4ef7-542d-bbb7-c09d3a79fcae"
 version = "0.8.6"
 
 [[deps.Downloads]]
-deps = ["ArgTools", "LibCURL", "NetworkOptions"]
+deps = ["ArgTools", "FileWatching", "LibCURL", "NetworkOptions"]
 uuid = "f43a241f-c20a-4ad4-852c-f6b1247861c6"
 
+[[deps.ExprTools]]
+git-tree-sha1 = "56559bbef6ca5ea0c0818fa5c90320398a6fbf8d"
+uuid = "e2ba6199-217a-4e67-a87a-7c52f15ade04"
+version = "0.1.8"
+
+[[deps.FastClosures]]
+git-tree-sha1 = "acebe244d53ee1b461970f8910c235b259e772ef"
+uuid = "9aa1b823-49e4-5ca5-8b0f-3971ec8bab6a"
+version = "0.3.2"
+
+[[deps.FileWatching]]
+uuid = "7b1f6079-737a-58dc-b8bc-7a2ca5c1b5ee"
+
 [[deps.FillArrays]]
 deps = ["LinearAlgebra", "Random", "SparseArrays", "Statistics"]
 git-tree-sha1 = "246621d23d1f43e3b9c368bf3b72b2331a27c286"
@@ -266,6 +291,12 @@ git-tree-sha1 = "1169632f425f79429f245113b775a0e3d121457c"
 uuid = "42fd0dbc-a981-5370-80f2-aaf504508153"
 version = "0.9.2"
 
+[[deps.JET]]
+deps = ["InteractiveUtils", "JuliaInterpreter", "LoweredCodeUtils", "MacroTools", "Pkg", "Revise", "Test"]
+git-tree-sha1 = "8e78b0c297cfa6cefd579f87232c89bd6ed7a081"
+uuid = "c3a54625-cd67-489e-a8e7-0a5a0ff4e31b"
+version = "0.5.16"
+
 [[deps.JLLWrappers]]
 deps = ["Preferences"]
 git-tree-sha1 = "abc9885a7ca2052a736a600f7fa66209f96506e1"
@@ -290,12 +321,24 @@ git-tree-sha1 = "57c17a221a55f81890aabf00f478886859e25eaf"
 uuid = "4076af6c-e467-56ae-b986-b466b2749572"
 version = "0.21.5"
 
+[[deps.JuliaInterpreter]]
+deps = ["CodeTracking", "InteractiveUtils", "Random", "UUIDs"]
+git-tree-sha1 = "52617c41d2761cc05ed81fe779804d3b7f14fff7"
+uuid = "aa1ae85d-cabe-5617-a682-6adf51b2e16a"
+version = "0.9.13"
+
 [[deps.Krylov]]
 deps = ["LinearAlgebra", "Printf", "SparseArrays"]
 git-tree-sha1 = "82f5afb342a5624dc4651981584a841f6088166b"
 uuid = "ba0b0d4f-ebba-5204-a429-3ac8c609bfb7"
 version = "0.8.0"
 
+[[deps.LDLFactorizations]]
+deps = ["AMD", "LinearAlgebra", "SparseArrays", "Test"]
+git-tree-sha1 = "736e01b9b2d443c4e3351aebe551b8a374ab9c05"
+uuid = "40e66cde-538c-5869-a4ad-c39174c6795b"
+version = "0.8.2"
+
 [[deps.LibCURL]]
 deps = ["LibCURL_jll", "MozillaCACerts_jll"]
 uuid = "b27032c2-a3e7-50c8-80cd-2d36dbcbfd21"
@@ -325,6 +368,12 @@ version = "7.1.1"
 deps = ["Libdl", "libblastrampoline_jll"]
 uuid = "37e2e46d-f89d-539d-b4ee-838fcccc9c8e"
 
+[[deps.LinearOperators]]
+deps = ["FastClosures", "LDLFactorizations", "LinearAlgebra", "Printf", "SparseArrays", "TimerOutputs"]
+git-tree-sha1 = "b404faa9b85e62c0eeec7a600d5b4316c58215ed"
+uuid = "5c8ed15e-5a4c-59e4-a42b-c7e8811fb125"
+version = "2.3.2"
+
 [[deps.LogExpFunctions]]
 deps = ["ChainRulesCore", "ChangesOfVariables", "DocStringExtensions", "InverseFunctions", "IrrationalConstants", "LinearAlgebra"]
 git-tree-sha1 = "a970d55c2ad8084ca317a4658ba6ce99b7523571"
@@ -334,6 +383,12 @@ version = "0.3.12"
 [[deps.Logging]]
 uuid = "56ddb016-857b-54e1-b83d-db4d58db5568"
 
+[[deps.LoweredCodeUtils]]
+deps = ["JuliaInterpreter"]
+git-tree-sha1 = "dedbebe234e06e1ddad435f5c6f4b85cd8ce55f7"
+uuid = "6f1432cf-f94c-5a45-995e-cdbf5db27b0b"
+version = "2.2.2"
+
 [[deps.METIS_jll]]
 deps = ["Artifacts", "JLLWrappers", "Libdl", "Pkg"]
 git-tree-sha1 = "1d31872bb9c5e7ec1f618e8c4a56c8b0d9bddc7e"
@@ -531,6 +586,12 @@ git-tree-sha1 = "838a3a4188e2ded87a4f9f184b4b0d78a1e91cb7"
 uuid = "ae029012-a4dd-5104-9daa-d747884805df"
 version = "1.3.0"
 
+[[deps.Revise]]
+deps = ["CodeTracking", "Distributed", "FileWatching", "JuliaInterpreter", "LibGit2", "LoweredCodeUtils", "OrderedCollections", "Pkg", "REPL", "Requires", "UUIDs", "Unicode"]
+git-tree-sha1 = "4d4239e93531ac3e7ca7e339f15978d0b5149d03"
+uuid = "295af30f-e4ad-537b-8983-00126c2a3abe"
+version = "3.3.3"
+
 [[deps.Richardson]]
 deps = ["LinearAlgebra"]
 git-tree-sha1 = "e03ca566bec93f8a3aeb059c8ef102f268a38949"
@@ -628,6 +689,12 @@ uuid = "a4e569a6-e804-4fa4-b0f3-eef7a1d5b13e"
 deps = ["InteractiveUtils", "Logging", "Random", "Serialization"]
 uuid = "8dfed614-e22c-5e08-85e1-65c5234f0b40"
 
+[[deps.TimerOutputs]]
+deps = ["ExprTools", "Printf"]
+git-tree-sha1 = "7638550aaea1c9a1e86817a231ef0faa9aca79bd"
+uuid = "a759f4b9-e2f1-59dc-863e-4aeb61b1ea8f"
+version = "0.5.19"
+
 [[deps.TranscodingStreams]]
 deps = ["Random", "Test"]
 git-tree-sha1 = "216b95ea110b5972db65aa90f88d8d89dcb8851c"