Merge pull request #10 from gdalle/dev-giom

Prepare first release

Author: gdalle

Project.toml

@@ -5,9 +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

@@ -5,7 +5,9 @@
 [![Build Status](https://github.com/gdalle/ImplicitDifferentiation.jl/actions/workflows/CI.yml/badge.svg?branch=main)](https://github.com/gdalle/ImplicitDifferentiation.jl/actions/workflows/CI.yml?query=branch%3Amain)
 [![Coverage](https://codecov.io/gh/gdalle/ImplicitDifferentiation.jl/branch/main/graph/badge.svg)](https://codecov.io/gh/gdalle/ImplicitDifferentiation.jl)
 [![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!
+> This package is in a very early development stage, so use it at your own risk!

docs/Manifest.toml

@@ -123,6 +123,12 @@ version = "3.43.0"
 deps = ["Artifacts", "Libdl"]
 uuid = "e66e0078-7015-5450-92f7-15fbd957f2ae"
 
+[[deps.ComponentArrays]]
+deps = ["ArrayInterface", "ChainRulesCore", "LinearAlgebra", "Requires"]
+git-tree-sha1 = "243d8b8afc829a6707bbb1cd00da868703c2ef42"
+uuid = "b0b7db55-cfe3-40fc-9ded-d10e2dbeff66"
+version = "0.11.15"
+
 [[deps.DataAPI]]
 git-tree-sha1 = "cc70b17275652eb47bc9e5f81635981f13cea5c8"
 uuid = "9a962f9c-6df0-11e9-0e5d-c546b8b5ee8a"
@@ -192,12 +198,6 @@ version = "0.27.15"
 deps = ["ArgTools", "LibCURL", "NetworkOptions"]
 uuid = "f43a241f-c20a-4ad4-852c-f6b1247861c6"
 
-[[deps.ExactOptimalTransport]]
-deps = ["Distances", "Distributions", "LinearAlgebra", "MathOptInterface", "PDMats", "QuadGK", "SparseArrays", "StatsBase"]
-git-tree-sha1 = "fb24e2b311008117679935b4e3d43cf285359153"
-uuid = "24df6009-d856-477c-ac5c-91f668376b31"
-version = "0.1.1"
-
 [[deps.ExprTools]]
 git-tree-sha1 = "56559bbef6ca5ea0c0818fa5c90320398a6fbf8d"
 uuid = "e2ba6199-217a-4e67-a87a-7c52f15ade04"
@@ -256,7 +256,7 @@ uuid = "615f187c-cbe4-4ef1-ba3b-2fcf58d6d173"
 version = "0.1.1"
 
 [[deps.ImplicitDifferentiation]]
-deps = ["ChainRulesCore", "LinearOperators"]
+deps = ["ChainRulesCore", "LinearOperators", "SparseArrays"]
 path = ".."
 uuid = "57b37032-215b-411a-8a7c-41a003a55207"
 version = "0.1.0"
@@ -455,9 +455,9 @@ version = "7.8.2"
 
 [[deps.NNlib]]
 deps = ["Adapt", "ChainRulesCore", "Compat", "LinearAlgebra", "Pkg", "Requires", "Statistics"]
-git-tree-sha1 = "a59a614b8b4ea6dc1dcec8c6514e251f13ccbe10"
+git-tree-sha1 = "3a8dfd0cfb5bb3b82d09949e14423409b9334acb"
 uuid = "872c559c-99b0-510c-b3b7-b6c96a88d5cd"
-version = "0.8.4"
+version = "0.7.34"
 
 [[deps.NaNMath]]
 git-tree-sha1 = "b086b7ea07f8e38cf122f5016af580881ac914fe"
@@ -494,10 +494,10 @@ uuid = "429524aa-4258-5aef-a3af-852621145aeb"
 version = "1.6.2"
 
 [[deps.OptimalTransport]]
-deps = ["ExactOptimalTransport", "IterativeSolvers", "LinearAlgebra", "LogExpFunctions", "NNlib", "Reexport"]
-git-tree-sha1 = "79ba1dab46dfc7b677278ebe892a431788da86a9"
+deps = ["Distances", "Distributions", "IterativeSolvers", "LinearAlgebra", "LogExpFunctions", "MathOptInterface", "NNlib", "PDMats", "QuadGK", "SparseArrays", "StatsBase"]
+git-tree-sha1 = "d85a74ab73bb0f4ccec7f713a352ba4c11daf750"
 uuid = "7e02d93a-ae51-4f58-b602-d97af76e3b33"
-version = "0.3.19"
+version = "0.3.15"
 
 [[deps.OrderedCollections]]
 git-tree-sha1 = "85f8e6578bf1f9ee0d11e7bb1b1456435479d47c"

docs/Project.toml

@@ -1,8 +1,10 @@
 [deps]
 ChainRulesCore = "d360d2e6-b24c-11e9-a2a3-2a2ae2dbcce4"
 ChainRulesTestUtils = "cdddcdb0-9152-4a09-a978-84456f9df70a"
+ComponentArrays = "b0b7db55-cfe3-40fc-9ded-d10e2dbeff66"
 Distances = "b4f34e82-e78d-54a5-968a-f98e89d6e8f7"
 Documenter = "e30172f5-a6a5-5a46-863b-614d45cd2de4"
+FillArrays = "1a297f60-69ca-5386-bcde-b61e274b549b"
 FiniteDiff = "6a86dc24-6348-571c-b903-95158fe2bd41"
 ImplicitDifferentiation = "57b37032-215b-411a-8a7c-41a003a55207"
 Ipopt = "b6b21f68-93f8-5de0-b562-5493be1d77c9"

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,10 +1,14 @@
 module ImplicitDifferentiation
 
 using ChainRulesCore
+using Krylov
 using LinearOperators
+using SparseArrays
 
 include("implicit_function.jl")
+include("simplex.jl")
 
 export ImplicitFunction
+export simplex_projection
 
 end

src/implicit_function.jl

@@ -14,18 +14,34 @@ 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
+
 """
     implicit(x)
 
 Make [`ImplicitFunction`](@ref) callable by applying `implicit.forward`.
 """
-(implicit::ImplicitFunction)(x::AbstractVector{<:Real}) = implicit.forward(x)
+(implicit::ImplicitFunction)(x) = implicit.forward(x)
 
 """
     frule(rc, (_, dx), implicit, x)
@@ -34,30 +50,33 @@ Custom forward rule for [`ImplicitFunction`](@ref).
 
 We compute the Jacobian-vector product `Jv` by solving `Au = Bv` and setting `Jv = u`.
 """
-function ChainRulesCore.frule(rc::RuleConfig, (_, dx), implicit::ImplicitFunction, x::AbstractVector{R}) where {R<:Real}
-    forward = implicit.forward
-    conditions = implicit.conditions
-    linear_solver = implicit.linear_solver
+function ChainRulesCore.frule(
+    rc::RuleConfig, (_, dx), implicit::ImplicitFunction, x::AbstractVector
+)
+    (; forward, conditions, linear_solver) = implicit
 
     y = forward(x)
+    n, m = length(x), length(y)
 
-    F₁(x̃) = conditions(x̃, y)
-    F₂(ỹ) = -conditions(x, ỹ)
+    conditions_x(x̃) = conditions(x̃, y)
+    conditions_y(ỹ) = -conditions(x, ỹ)
 
-    pushforward_A(dỹ) = frule_via_ad(rc, (NoTangent(), dỹ), F₂, y)[2]
-    pushforward_B(dx̃) = frule_via_ad(rc, (NoTangent(), dx̃), F₁, x)[2]
+    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]
 
-    n, m, c = length(x), length(y), length(y)
     mul_A!(res, v) = res .= pushforward_A(v)
     mul_B!(res, v) = res .= pushforward_B(v)
 
-    A = LinearOperator(R, c, m, false, false, mul_A!)
-    B = LinearOperator(R, c, m, false, false, mul_B!)
+    A = LinearOperator(Float64, m, m, false, false, mul_A!)
+    B = LinearOperator(Float64, m, n, false, false, mul_B!)
 
-    b = B * unthunk(dx)
-    dy, stats = linear_solver(A, b)
-
-    return y, dy
+    dx_vec = Vector(unthunk(dx))
+    b = B * dx_vec
+    dy_vec, stats = linear_solver(A, b)
+    if !stats.solved
+        throw(SolverFailureException("The linear solver failed to converge"))
+    end
+    return y, dy_vec
 end
 
 """
@@ -67,30 +86,32 @@ Custom reverse rule for [`ImplicitFunction`](@ref).
 
 We compute the vector-Jacobian product `Jᵀv` by solving `Aᵀu = v` and setting `Jᵀv = Bᵀu`.
 """
-function ChainRulesCore.rrule(rc::RuleConfig, implicit::ImplicitFunction, x::AbstractVector{R}) where {R<:Real}
-    forward = implicit.forward
-    conditions = implicit.conditions
-    linear_solver = implicit.linear_solver
+function ChainRulesCore.rrule(rc::RuleConfig, implicit::ImplicitFunction, x::AbstractVector)
+    (; forward, conditions, linear_solver) = implicit
 
     y = forward(x)
+    n, m = length(x), length(y)
 
-    F₁(x̃) = conditions(x̃, y)
-    F₂(ỹ) = -conditions(x, ỹ)
+    conditions_x(x̃) = conditions(x̃, y)
+    conditions_y(ỹ) = -conditions(x, ỹ)
 
-    pullback_Aᵀ = last ∘ rrule_via_ad(rc, F₂, y)[2]
-    pullback_Bᵀ = last ∘ rrule_via_ad(rc, F₁, x)[2]
+    pullback_Aᵀ = last ∘ rrule_via_ad(rc, conditions_y, y)[2]
+    pullback_Bᵀ = last ∘ rrule_via_ad(rc, conditions_x, x)[2]
 
-    n, m, c = length(x), length(y), length(y)
     mul_Aᵀ!(res, v) = res .= pullback_Aᵀ(v)
     mul_Bᵀ!(res, v) = res .= pullback_Bᵀ(v)
 
-    Aᵀ = LinearOperator(R, m, c, false, false, mul_Aᵀ!)
-    Bᵀ = LinearOperator(R, n, c, false, false, mul_Bᵀ!)
+    Aᵀ = LinearOperator(Float64, m, m, false, false, mul_Aᵀ!)
+    Bᵀ = LinearOperator(Float64, n, m, false, false, mul_Bᵀ!)
 
     function implicit_pullback(dy)
-        u, stats = linear_solver(Aᵀ, unthunk(dy))
-        dx = Bᵀ * u
-        return (NoTangent(), dx)
+        dy_vec = Vector(unthunk(dy))
+        u, stats = linear_solver(Aᵀ, dy_vec)
+        if !stats.solved
+            throw(SolverFailureException("The linear solver failed to converge"))
+        end
+        dx_vec = Bᵀ * u
+        return (NoTangent(), dx_vec)
     end
 
     return y, implicit_pullback

src/simplex.jl

@@ -0,0 +1,48 @@
+"""
+    simplex_projection_and_support(z)
+
+Compute the Euclidean projection onto the probability simplex and the set of indices where it is nonzero.
+
+See <https://arxiv.org/abs/1602.02068> for details.
+"""
+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 = 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;
+
+"""
+    simplex_projection(z)
+
+Compute the Euclidean projection onto the probability simplex.
+"""
+function simplex_projection(z::AbstractVector{<:Real})
+    p, _ = simplex_projection_and_support(z)
+    return p
+end;
+
+"""
+    rrule(::typeof(simplex_projection), z)
+
+Custom reverse rule for [`simplex_projection`](@ref) which bypasses the sorting step.
+
+See <https://arxiv.org/abs/1602.02068> for details.
+"""
+function ChainRulesCore.rrule(::typeof(simplex_projection), z::AbstractVector{<:Real})
+    p, s = simplex_projection_and_support(z)
+    S = sum(s)
+    function simplex_projection_pullback(dp)
+        vjp = s .* (dp .- (dp's) / S)
+        return (NoTangent(), vjp)
+    end
+    return p, simplex_projection_pullback
+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
 
@@ -46,7 +45,7 @@ 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
 
@@ -65,6 +64,6 @@ Zygote.jacobian(implicit, x)[1]
 # The following tests are not included in the docs.  #src
 
 @testset verbose = true "ChainRules" begin  #src
-    test_frule(implicit, x)  #src
-    test_rrule(implicit, x)  #src
+    test_frule(implicit, x; check_inferred=false)  #src
+    test_rrule(implicit, x; check_inferred=false)  #src
 end  #src