Update to DI v0.2 (#139)

Author: gdalle

.github/workflows/CI.yml

@@ -50,7 +50,7 @@ jobs:
           version: '1'
       - uses: julia-actions/cache@v1
       - name: Install dependencies
-        run: julia --project=docs/ -e 'using Pkg; Pkg.develop(PackageSpec(path=pwd())); Pkg.instantiate()'
+        run: julia --project=docs/ -e 'using Pkg; Pkg.Registry.update(); Pkg.develop(PackageSpec(path=pwd())); Pkg.instantiate()'
       - name: Build and deploy
         env:
           GITHUB_TOKEN: ${{ secrets.GITHUB_TOKEN }} # If authenticating with GitHub Actions token

Project.toml

@@ -23,7 +23,7 @@ ImplicitDifferentiationForwardDiffExt = "ForwardDiff"
 [compat]
 ADTypes = "0.2"
 ChainRulesCore = "1.23.0"
-DifferentiationInterface = "0.1"
+DifferentiationInterface = "0.2"
 Enzyme = "0.11.20"
 ForwardDiff = "0.10.36"
 Krylov = "0.9.5"

docs/src/faq.md

@@ -2,7 +2,7 @@
 
 ## Supported autodiff backends
 
-To differentiate an `ImplicitFunction`, the following backends are supported.
+To differentiate through an `ImplicitFunction`, the following backends are supported.
 
 | Backend                                                                | Forward mode | Reverse mode |
 | :--------------------------------------------------------------------- | :----------- | :----------- |
@@ -31,9 +31,9 @@ Or better yet, wrap it in a static vector: `SVector(val)`.
 ### Sparse arrays
 
 !!! danger "Danger"
-    Sparse arrays are not officially supported and might give incorrect values or `NaN`s!
+    Sparse arrays are not supported and might give incorrect values or `NaN`s!
 
-With ForwardDiff.jl, differentiation of sparse arrays will always give wrong results due to [sparsity pattern cancellation](https://github.com/JuliaDiff/ForwardDiff.jl/issues/658).
+With ForwardDiff.jl, differentiation of sparse arrays will often give wrong results due to [sparsity pattern cancellation](https://github.com/JuliaDiff/ForwardDiff.jl/issues/658).
 That is why we do not test behavior for sparse inputs.
 
 ## Number of inputs and outputs
@@ -92,6 +92,21 @@ This is mainly useful when the solution procedure creates objects such as Jacobi
 In that case, you may want to write the conditions differentiation rules yourself.
 A more advanced application is given by [DifferentiableFrankWolfe.jl](https://github.com/gdalle/DifferentiableFrankWolfe.jl).
 
+## Linear system
+
+### Lazy or dense
+
+Usually, dense Jacobians are more efficient in small dimension, while lazy operators become necessary in high dimension.
+This choice is made via the `lazy` type parameter of [`ImplicitFunction`](@ref), with `lazy = true` being the default.
+
+### Picking a solver
+
+The right linear solver to use depends on the Jacobian representation.
+You can usually stick to the default settings:
+
+- the direct solver `\` for dense Jacobians
+- an iterative solver from [Krylov.jl](https://github.com/JuliaSmoothOptimizers/Krylov.jl) for lazy operators
+
 ## Modeling tips
 
 ### Writing conditions

examples/1_basic.jl

@@ -73,7 +73,7 @@ end;
 We now have all the ingredients to construct our implicit function.
 =#
 
-implicit_optim = ImplicitFunction(; forward=forward_optim, conditions=conditions_optim)
+implicit_optim = ImplicitFunction(forward_optim, conditions_optim)
 
 # And indeed, it behaves as it should when we call it:
 

src/ImplicitDifferentiation.jl

@@ -9,11 +9,7 @@ module ImplicitDifferentiation
 
 using ADTypes: AbstractADType
 using DifferentiationInterface:
-    jacobian,
-    prepare_pushforward,
-    prepare_pullback,
-    pushforward!!,
-    value_and_pullback!!_split
+    jacobian, prepare_pushforward, prepare_pullback, pushforward!, value_and_pullback!_split
 using Krylov: block_gmres, gmres
 using LinearOperators: LinearOperator
 using LinearAlgebra: factorize, lu

src/implicit_function.jl

@@ -14,7 +14,7 @@ function (::DefaultLinearSolver)(A, B::AbstractMatrix)
 end
 
 """
-    ImplicitFunction
+    ImplicitFunction{lazy}
 
 Wrapper for an implicit function defined by a forward mapping `y` and a set of conditions `c`.
 
@@ -27,47 +27,94 @@ This requires solving a linear system `A * J = -B` where `A = ∂c/∂y`, `B = 
 
 # Fields
 
-- `forward`: a callable, does not need to be compatible with automatic differentiation
-- `conditions`: a callable, must be compatible with automatic differentiation
-- `linear_solver`: a callable with two methods:
-    - `(A, b::AbstractVector) -> s::AbstractVector` such that `A * s = b`
-    - `(A, B::AbstractVector) -> S::AbstractMatrix` such that `A * S = B` 
-- `conditions_x_backend`: either `nothing` or an object subtyping `AbstractADType` from [ADTypes.jl](https://github.com/SciML/ADTypes.jl), defines how the conditions will be differentiated with respect to the first argument `x` 
-- `conditions_y_backend`: same for the second argument `y`
+- `forward`: a callable computing `y(x)`, does not need to be compatible with automatic differentiation
+- `conditions`: a callable computing `c(x, y)`, must be compatible with automatic differentiation
+- `linear_solver`: a callable to solve the linear system `A * J = -B`
+- `conditions_x_backend`: defines how the conditions will be differentiated with respect to the first argument `x` 
+- `conditions_y_backend`: defines how the conditions will be differentiated with respect to the second argument `y` 
+
+# Type parameters
+
+- `lazy`: whether to use a `LinearOperator` from [LinearOperators.jl](https://github.com/JuliaSmoothOptimizers/LinearOperators.jl) (`lazy = true`) or a dense Jacobian matrix (`lazy = false`) for `A` and `B`
+
+# Constructors
+
+    ImplicitFunction{lazy}(
+        forward, conditions;
+        linear_solver, conditions_x_backend, conditions_x_backend
+    )
+
+    ImplicitFunction(
+        forward, conditions;
+        linear_solver, conditions_x_backend, conditions_x_backend
+    )
+
+Default values:
+
+- `lazy = true`
+- `linear_solver`: the direct solver `\` for dense Jacobians, or an iterative solver from [Krylov.jl](https://github.com/JuliaSmoothOptimizers/Krylov.jl) for lazy operators
+- `conditions_x_backend = nothing`
+- `conditions_y_backend = nothing`
+
+# Function signatures
 
 There are two possible signatures for `forward` and `conditions`, which must be consistent with one another:
     
 | standard | byproduct |
 |:---|:---|
-| `forward(x, args...; kwargs...) = y` | `conditions(x, y, args...; kwargs...) = c` |
+| `forward(x, args...; kwargs...) = y`      | `conditions(x, y, args...; kwargs...) = c` |
 | `forward(x, args...; kwargs...) = (y, z)` | `conditions(x, y, z, args...; kwargs...) = c` |
     
 In both cases, `x`, `y` and `c` must be `AbstractVector`s, with `length(y) = length(c)`.
 In the second case, the byproduct `z` can be an arbitrary object generated by `forward`.
 The positional arguments `args...` and keyword arguments `kwargs...` must be the same for both `forward` and `conditions`.
 
 The byproduct `z` and the other positional arguments `args...` beyond `x` are considered constant for differentiation purposes.
+
+# Linear solver
+
+The provided `linear_solver` objects needs to be callable, with two methods:
+- `(A, b::AbstractVector) -> s::AbstractVector` such that `A * s = b`
+- `(A, B::AbstractVector) -> S::AbstractMatrix` such that `A * S = B` 
+
+# Condition backends
+
+The provided `conditions_x_backend` and `conditions_y_backend` can be either:
+- `nothing`, in which case the outer backend (the one differentiating through the `ImplicitFunction`) is used
+- an object subtyping `AbstractADType` from [ADTypes.jl](https://github.com/SciML/ADTypes.jl).
 """
-@kwdef struct ImplicitFunction{
-    F,C,L,B1<:Union{Nothing,AbstractADType},B2<:Union{Nothing,AbstractADType}
+struct ImplicitFunction{
+    lazy,F,C,L,B1<:Union{Nothing,AbstractADType},B2<:Union{Nothing,AbstractADType}
 }
     forward::F
     conditions::C
-    linear_solver::L = DefaultLinearSolver()
-    conditions_x_backend::B1 = nothing
-    conditions_y_backend::B2 = nothing
+    linear_solver::L
+    conditions_x_backend::B1
+    conditions_y_backend::B2
+end
+
+function ImplicitFunction{lazy}(
+    forward::F,
+    conditions::C;
+    linear_solver::L=lazy ? DefaultLinearSolver() : \,
+    conditions_x_backend::B1=nothing,
+    conditions_y_backend::B2=nothing,
+) where {lazy,F,C,L,B1,B2}
+    return ImplicitFunction{lazy,F,C,L,B1,B2}(
+        forward, conditions, linear_solver, conditions_x_backend, conditions_y_backend
+    )
 end
 
 function ImplicitFunction(forward, conditions; kwargs...)
-    return ImplicitFunction(; forward, conditions, kwargs...)
+    return ImplicitFunction{true}(forward, conditions; kwargs...)
 end
 
-function Base.show(io::IO, implicit::ImplicitFunction)
+function Base.show(io::IO, implicit::ImplicitFunction{lazy}) where {lazy}
     (; forward, conditions, linear_solver, conditions_x_backend, conditions_y_backend) =
         implicit
     return print(
         io,
-        "ImplicitFunction($forward, $conditions, $linear_solver, $conditions_x_backend, $conditions_y_backend)",
+        "ImplicitFunction{$(lazy ? "lazy" : "dense")}($forward, $conditions, $linear_solver, $conditions_x_backend, $conditions_y_backend)",
     )
 end
 

src/operators.jl

@@ -81,51 +81,47 @@ end
 
 function (po::PushforwardOperator!)(res, v, α, β)
     if iszero(β)
-        res .= pushforward!!(po.f, res, po.backend, po.x, v, po.extras)
+        pushforward!(po.f, res, po.backend, po.x, v, po.extras)
         res .= α .* res
     else
         po.res_backup .= res
-        res .= pushforward!!(po.f, res, po.backend, po.x, v, po.extras)
+        pushforward!(po.f, res, po.backend, po.x, v, po.extras)
         res .= α .* res .+ β .* po.res_backup
     end
     return res
 end
 
 struct PullbackOperator!{PB,R}
-    pullbackfunc!!::PB
+    pullbackfunc!::PB
     res_backup::R
 end
 
 function (po::PullbackOperator!)(res, v, α, β)
     if iszero(β)
-        res .= po.pullbackfunc!!(res, v)
-        res .= α .* res
+        po.pullbackfunc!(res, v)
     else
         po.res_backup .= res
-        res .= po.pullbackfunc!!(res, v)
+        po.pullbackfunc!(res, v)
         res .= α .* res .+ β .+ po.res_backup
     end
     return res
 end
 
 function build_A(
-    implicit::ImplicitFunction,
+    implicit::ImplicitFunction{lazy},
     x::AbstractVector,
     y_or_yz,
     args...;
     suggested_backend,
     kwargs...,
-)
+) where {lazy}
     (; conditions, linear_solver, conditions_y_backend) = implicit
     y = output(y_or_yz)
     n, m = length(x), length(y)
     back_y = isnothing(conditions_y_backend) ? suggested_backend : conditions_y_backend
     cond_y = ConditionsY(conditions, x, y_or_yz, args...; kwargs...)
-    if linear_solver isa typeof(\)
-        J = jacobian(cond_y, back_y, y)
-        A = factorize(J)
-    else
-        extras = prepare_pushforward(cond_y, back_y, y)
+    if lazy
+        extras = prepare_pushforward(cond_y, back_y, y, similar(y))
         A = LinearOperator(
             eltype(y),
             m,
@@ -135,59 +131,60 @@ function build_A(
             PushforwardOperator!(cond_y, back_y, y, extras, similar(y)),
             typeof(y),
         )
+    else
+        J = jacobian(cond_y, back_y, y)
+        A = factorize(J)
     end
     return A
 end
 
 function build_Aᵀ(
-    implicit::ImplicitFunction,
+    implicit::ImplicitFunction{lazy},
     x::AbstractVector,
     y_or_yz,
     args...;
     suggested_backend,
     kwargs...,
-)
+) where {lazy}
     (; conditions, linear_solver, conditions_y_backend) = implicit
     y = output(y_or_yz)
     n, m = length(x), length(y)
     back_y = isnothing(conditions_y_backend) ? suggested_backend : conditions_y_backend
     cond_y = ConditionsY(conditions, x, y_or_yz, args...; kwargs...)
-    if linear_solver isa typeof(\)
-        Jᵀ = transpose(jacobian(cond_y, back_y, y))
-        Aᵀ = factorize(Jᵀ)
-    else
-        extras = prepare_pullback(cond_y, back_y, y)
-        _, pullbackfunc!! = value_and_pullback!!_split(cond_y, back_y, y, extras)
+    if lazy
+        extras = prepare_pullback(cond_y, back_y, y, similar(y))
+        _, pullbackfunc! = value_and_pullback!_split(cond_y, back_y, y, extras)
         Aᵀ = LinearOperator(
             eltype(y),
             m,
             m,
             false,
             false,
-            PullbackOperator!(pullbackfunc!!, similar(y)),
+            PullbackOperator!(pullbackfunc!, similar(y)),
             typeof(y),
         )
+    else
+        Jᵀ = transpose(jacobian(cond_y, back_y, y))
+        Aᵀ = factorize(Jᵀ)
     end
     return Aᵀ
 end
 
 function build_B(
-    implicit::ImplicitFunction,
+    implicit::ImplicitFunction{lazy},
     x::AbstractVector,
     y_or_yz,
     args...;
     suggested_backend,
     kwargs...,
-)
+) where {lazy}
     (; conditions, linear_solver, conditions_x_backend) = implicit
     y = output(y_or_yz)
     n, m = length(x), length(y)
     back_x = isnothing(conditions_x_backend) ? suggested_backend : conditions_x_backend
     cond_x = ConditionsX(conditions, x, y_or_yz, args...; kwargs...)
-    if linear_solver isa typeof(\)
-        B = transpose(jacobian(cond_x, back_x, x))
-    else
-        extras = prepare_pushforward(cond_x, back_x, x)
+    if lazy
+        extras = prepare_pushforward(cond_x, back_x, x, similar(x))
         B = LinearOperator(
             eltype(y),
             m,
@@ -197,37 +194,39 @@ function build_B(
             PushforwardOperator!(cond_x, back_x, x, extras, similar(y)),
             typeof(x),
         )
+    else
+        B = transpose(jacobian(cond_x, back_x, x))
     end
     return B
 end
 
 function build_Bᵀ(
-    implicit::ImplicitFunction,
+    implicit::ImplicitFunction{lazy},
     x::AbstractVector,
     y_or_yz,
     args...;
     suggested_backend,
     kwargs...,
-)
+) where {lazy}
     (; conditions, linear_solver, conditions_x_backend) = implicit
     y = output(y_or_yz)
     n, m = length(x), length(y)
     back_x = isnothing(conditions_x_backend) ? suggested_backend : conditions_x_backend
     cond_x = ConditionsX(conditions, x, y_or_yz, args...; kwargs...)
-    if linear_solver isa typeof(\)
-        Bᵀ = transpose(jacobian(cond_x, back_x, x))
-    else
-        extras = prepare_pullback(cond_x, back_x, x)
-        _, pullbackfunc!! = value_and_pullback!!_split(cond_x, back_x, x, extras)
+    if lazy
+        extras = prepare_pullback(cond_x, back_x, x, similar(y))
+        _, pullbackfunc! = value_and_pullback!_split(cond_x, back_x, x, extras)
         Bᵀ = LinearOperator(
             eltype(y),
             n,
             m,
             false,
             false,
-            PullbackOperator!(pullbackfunc!!, similar(y)),
+            PullbackOperator!(pullbackfunc!, similar(y)),
             typeof(x),
         )
+    else
+        Bᵀ = transpose(jacobian(cond_x, back_x, x))
     end
     return Bᵀ
 end

test/runtests.jl

@@ -10,7 +10,7 @@ using Random
 using Test
 using Zygote: Zygote
 
-DocMeta.setdocmeta!(
+Documenter.DocMeta.setdocmeta!(
     ImplicitDifferentiation, :DocTestSetup, :(using ImplicitDifferentiation); recursive=true
 )