Split out preparation

Author: gdalle

docs/src/faq.md

@@ -18,8 +18,7 @@ However, this can be switched to any other "inner" backend compatible with [Diff
 ### Arrays
 
 Functions that eat or spit out arbitrary arrays are supported, as long as the forward mapping _and_ conditions return arrays of the same size.
-
-If you deal with small arrays (say, less than 100 elements), consider using [StaticArrays.jl](https://github.com/JuliaArrays/StaticArrays.jl) for increased performance.
+The array types involved should be mutable.
 
 ### Scalars
 

examples/3_tricks.jl

@@ -38,21 +38,19 @@ end
 
 function conditions_components(x::ComponentVector, y::ComponentVector, _z)
     c_d, c_e = conditions_components_aux(x.a, x.b, x.m, y.d, y.e)
-    c = ComponentVector(; c_d=c_d, c_e=c_e)
+    c = ComponentVector(; d=c_d, e=c_e)
     return c
 end;
 
 # And build your implicit function like so:
 
-implicit_components = ImplicitFunction(
-    forward_components, conditions_components; strict=Val(false)
-);
+implicit_components = ImplicitFunction(forward_components, conditions_components);
 
 # Now we're good to go.
 
 a, b, m = [1.0, 2.0], [3.0, 4.0, 5.0], 6.0
 x = ComponentVector(; a=a, b=b, m=m)
-implicit_components(x)
+y, z = implicit_components(x)
 
 # And it works with both ForwardDiff.jl and Zygote.jl
 

ext/ImplicitDifferentiationChainRulesCoreExt.jl

@@ -6,6 +6,7 @@ using ChainRulesCore: unthunk, @not_implemented
 using ImplicitDifferentiation:
     ImplicitDifferentiation,
     ImplicitFunction,
+    ImplicitFunctionPreparation,
     build_Aᵀ,
     build_Bᵀ,
     chainrules_suggested_backend
@@ -14,29 +15,33 @@ using ImplicitDifferentiation:
 ImplicitDifferentiation.chainrules_suggested_backend(rc::RuleConfig) = AutoChainRules(rc)
 
 function ChainRulesCore.rrule(
-    rc::RuleConfig, implicit::ImplicitFunction, x::AbstractArray, args::Vararg{Any,N};
+    rc::RuleConfig,
+    implicit::ImplicitFunction,
+    prep::ImplicitFunctionPreparation,
+    x::AbstractArray,
+    args::Vararg{Any,N};
 ) where {N}
     y, z = implicit(x, args...)
     c = implicit.conditions(x, y, z, args...)
 
     suggested_backend = chainrules_suggested_backend(rc)
-    Aᵀ = build_Aᵀ(implicit, x, y, z, c, args...; suggested_backend)
-    Bᵀ = build_Bᵀ(implicit, x, y, z, c, args...; suggested_backend)
+    Aᵀ = build_Aᵀ(implicit, prep, x, y, z, c, args...; suggested_backend)
+    Bᵀ = build_Bᵀ(implicit, prep, x, y, z, c, args...; suggested_backend)
     project_x = ProjectTo(x)
 
-    function implicit_pullback((dy, dz))
+    function implicit_pullback_prepared((dy, dz))
         dy = unthunk(dy)
         dy_vec = vec(dy)
-        dc_vec = similar(vec(c))
-        implicit.linear_solver(dc_vec, Aᵀ, -dy_vec)
+        dc_vec = implicit.linear_solver(Aᵀ, -dy_vec)
         dx_vec = Bᵀ(dc_vec)
         dx = reshape(dx_vec, size(x))
         df = NoTangent()
+        dprep = @not_implemented("Tangents for mutable arguments are not defined")
         dargs = ntuple(unimplemented_tangent, N)
-        return (df, project_x(dx), dargs...)
+        return (df, dprep, project_x(dx), dargs...)
     end
 
-    return (y, z), implicit_pullback
+    return (y, z), implicit_pullback_prepared
 end
 
 function unimplemented_tangent(_)

ext/ImplicitDifferentiationForwardDiffExt.jl

@@ -2,18 +2,19 @@ module ImplicitDifferentiationForwardDiffExt
 
 using ADTypes: AutoForwardDiff
 using ForwardDiff: Dual, Partials, partials, value
-using ImplicitDifferentiation: ImplicitFunction, build_A, build_B
+using ImplicitDifferentiation:
+    ImplicitFunction, ImplicitFunctionPreparation, build_A, build_B
 
 function (implicit::ImplicitFunction)(
-    x_and_dx::AbstractArray{Dual{T,R,N}}, args...
+    prep::ImplicitFunctionPreparation, x_and_dx::AbstractArray{Dual{T,R,N}}, args...
 ) where {T,R,N}
     x = value.(x_and_dx)
     y, z = implicit(x, args...)
     c = implicit.conditions(x, y, z, args...)
 
     suggested_backend = AutoForwardDiff()
-    A = build_A(implicit, x, y, z, c, args...; suggested_backend)
-    B = build_B(implicit, x, y, z, c, args...; suggested_backend)
+    A = build_A(implicit, prep, x, y, z, c, args...; suggested_backend)
+    B = build_B(implicit, prep, x, y, z, c, args...; suggested_backend)
 
     dX = ntuple(Val(N)) do k
         partials.(x_and_dx, k)
@@ -24,8 +25,7 @@ function (implicit::ImplicitFunction)(
         return dₖc_vec
     end
     dY = map(dC_vec) do dₖc_vec
-        dₖy_vec = similar(vec(y))
-        implicit.linear_solver(dₖy_vec, A, -dₖc_vec)
+        dₖy_vec = implicit.linear_solver(A, -dₖc_vec)
         dₖy = reshape(dₖy_vec, size(y))
         return dₖy
     end

src/ImplicitDifferentiation.jl

@@ -21,20 +21,21 @@ using DifferentiationInterface:
     pullback,
     pushforward!,
     pushforward
-using Krylov: KrylovConstructor, krylov_workspace, krylov_solve!, solution
-using LinearAlgebra: ldiv!
+using Krylov: Krylov, krylov_workspace, krylov_solve!, solution
 using LinearOperators: LinearOperator
 using LinearMaps: FunctionMap
 using LinearAlgebra: factorize
 
 include("utils.jl")
 include("settings.jl")
-include("preparation.jl")
 include("implicit_function.jl")
+include("preparation.jl")
 include("execution.jl")
+include("callable.jl")
 
 export MatrixRepresentation, OperatorRepresentation
 export IterativeLinearSolver, DirectLinearSolver
 export ImplicitFunction
+export prepare_implicit
 
 end

src/callable.jl

@@ -0,0 +1,9 @@
+function (implicit::ImplicitFunction)(x::AbstractArray, args::Vararg{Any,N}) where {N}
+    return implicit(ImplicitFunctionPreparation(), x, args...)
+end
+
+function (implicit::ImplicitFunction)(
+    ::ImplicitFunctionPreparation, x::AbstractArray, args::Vararg{Any,N}
+) where {N}
+    return implicit.solver(x, args...)
+end

src/execution.jl

@@ -1,35 +1,40 @@
-struct JVP!{F,P,B,I,C}
+struct JVP!{F,P,B,I,V,C}
     f::F
     prep::P
     backend::B
     input::I
+    v_buffer::V
     contexts::C
 end
 
-struct VJP!{F,P,B,I,C}
+struct VJP!{F,P,B,I,V,C}
     f::F
     prep::P
     backend::B
     input::I
+    v_buffer::V
     contexts::C
 end
 
-function (po::JVP!)(res::AbstractVector, v::AbstractVector)
-    (; f, backend, input, contexts, prep) = po
-    pushforward!(f, (res,), prep, backend, input, (v,), contexts...)
+function (po::JVP!)(res::AbstractVector, v_wrongtype::AbstractVector)
+    (; f, backend, input, v_buffer, contexts, prep) = po
+    copyto!(v_buffer, v_wrongtype)
+    pushforward!(f, (res,), prep, backend, input, (v_buffer,), contexts...)
     return res
 end
 
-function (po::VJP!)(res::AbstractVector, v::AbstractVector)
-    (; f, backend, input, contexts, prep) = po
-    pullback!(f, (res,), prep, backend, input, (v,), contexts...)
+function (po::VJP!)(res::AbstractVector, v_wrongtype::AbstractVector)
+    (; f, backend, input, v_buffer, contexts, prep) = po
+    copyto!(v_buffer, v_wrongtype)
+    pullback!(f, (res,), prep, backend, input, (v_buffer,), contexts...)
     return res
 end
 
 ## A
 
 function build_A(
     implicit::ImplicitFunction,
+    prep::ImplicitFunctionPreparation,
     x::AbstractArray,
     y::AbstractArray,
     z,
@@ -38,14 +43,15 @@ function build_A(
     suggested_backend::AbstractADType,
 )
     return build_A_aux(
-        implicit.representation, implicit, x, y, z, c, args...; suggested_backend
+        implicit.representation, implicit, prep, x, y, z, c, args...; suggested_backend
     )
 end
 
 function build_A_aux(
-    ::MatrixRepresentation, implicit, x, y, z, c, args...; suggested_backend
+    ::MatrixRepresentation, implicit, prep, x, y, z, c, args...; suggested_backend
 )
-    (; conditions, backends, prep_A) = implicit
+    (; conditions, backends) = implicit
+    (; prep_A) = prep
     actual_backend = isnothing(backends) ? suggested_backend : backends.y
     contexts = (Constant(x), Constant(z), map(Constant, args)...)
     if isnothing(prep_A)
@@ -59,6 +65,7 @@ end
 function build_A_aux(
     ::OperatorRepresentation{package,symmetric,hermitian,posdef},
     implicit,
+    prep,
     x,
     y,
     z,
@@ -67,7 +74,8 @@ function build_A_aux(
     suggested_backend,
 ) where {package,symmetric,hermitian,posdef}
     T = Base.promote_eltype(x, y, c)
-    (; conditions, backends, prep_A) = implicit
+    (; conditions, backends) = implicit
+    (; prep_A) = prep
     actual_backend = isnothing(backends) ? suggested_backend : backends.y
     contexts = (Constant(x), Constant(z), map(Constant, args)...)
     f_vec = VecToVec(Switch12(conditions), y)
@@ -82,9 +90,9 @@ function build_A_aux(
             f_vec, prep_A, actual_backend, y_vec, (dy_vec,), contexts...
         )
     end
-    prod! = JVP!(f_vec, prep_A_same, actual_backend, y_vec, contexts)
+    prod! = JVP!(f_vec, prep_A_same, actual_backend, y_vec, dy_vec, contexts)
     if package == :LinearOperators
-        return LinearOperator(T, length(c), length(y), symmetric, hermitian, prod!;)
+        return LinearOperator(T, length(c), length(y), symmetric, hermitian, prod!)
     elseif package == :LinearMaps
         return FunctionMap{T}(
             prod!,
@@ -102,6 +110,7 @@ end
 
 function build_Aᵀ(
     implicit::ImplicitFunction,
+    prep::ImplicitFunctionPreparation,
     x::AbstractArray,
     y::AbstractArray,
     z,
@@ -110,14 +119,15 @@ function build_Aᵀ(
     suggested_backend::AbstractADType,
 )
     return build_Aᵀ_aux(
-        implicit.representation, implicit, x, y, z, c, args...; suggested_backend
+        implicit.representation, implicit, prep, x, y, z, c, args...; suggested_backend
     )
 end
 
 function build_Aᵀ_aux(
-    ::MatrixRepresentation, implicit, x, y, z, c, args...; suggested_backend
+    ::MatrixRepresentation, implicit, prep, x, y, z, c, args...; suggested_backend
 )
-    (; conditions, backends, prep_Aᵀ) = implicit
+    (; conditions, backends) = implicit
+    (; prep_Aᵀ) = prep
     actual_backend = isnothing(backends) ? suggested_backend : backends.y
     contexts = (Constant(x), Constant(z), map(Constant, args)...)
     if isnothing(prep_Aᵀ)
@@ -133,6 +143,7 @@ end
 function build_Aᵀ_aux(
     ::OperatorRepresentation{package,symmetric,hermitian,posdef},
     implicit,
+    prep,
     x,
     y,
     z,
@@ -141,7 +152,8 @@ function build_Aᵀ_aux(
     suggested_backend,
 ) where {package,symmetric,hermitian,posdef}
     T = Base.promote_eltype(x, y, c)
-    (; conditions, backends, prep_Aᵀ) = implicit
+    (; conditions, backends) = implicit
+    (; prep_Aᵀ) = prep
     actual_backend = isnothing(backends) ? suggested_backend : backends.y
     contexts = (Constant(x), Constant(z), map(Constant, args)...)
     f_vec = VecToVec(Switch12(conditions), y)
@@ -156,9 +168,9 @@ function build_Aᵀ_aux(
             f_vec, prep_Aᵀ, actual_backend, y_vec, (dc_vec,), contexts...
         )
     end
-    prod! = VJP!(f_vec, prep_Aᵀ_same, actual_backend, y_vec, contexts)
+    prod! = VJP!(f_vec, prep_Aᵀ_same, actual_backend, y_vec, dc_vec, contexts)
     if package == :LinearOperators
-        return LinearOperator(T, length(y), length(c), symmetric, hermitian, prod!;)
+        return LinearOperator(T, length(y), length(c), symmetric, hermitian, prod!)
     elseif package == :LinearMaps
         return FunctionMap{T}(
             prod!,
@@ -176,14 +188,16 @@ end
 
 function build_B(
     implicit::ImplicitFunction,
+    prep::ImplicitFunctionPreparation,
     x::AbstractArray,
     y::AbstractArray,
     z,
     c,
     args...;
     suggested_backend::AbstractADType,
 )
-    (; conditions, backends, prep_B) = implicit
+    (; conditions, backends) = implicit
+    (; prep_B) = prep
     actual_backend = isnothing(backends) ? suggested_backend : backends.x
     contexts = (Constant(y), Constant(z), map(Constant, args)...)
     f_vec = VecToVec(conditions, x)
@@ -198,7 +212,8 @@ function build_B(
             f_vec, prep_B, actual_backend, x_vec, (dx_vec,), contexts...
         )
     end
-    function B_fun(dx_vec)
+    function B_fun(dx_vec_wrongtype)
+        copyto!(dx_vec, dx_vec_wrongtype)
         return pushforward(
             f_vec, prep_B_same, actual_backend, x_vec, (dx_vec,), contexts...
         )[1]
@@ -210,14 +225,16 @@ end
 
 function build_Bᵀ(
     implicit::ImplicitFunction,
+    prep::ImplicitFunctionPreparation,
     x::AbstractArray,
     y::AbstractArray,
     z,
     c,
     args...;
     suggested_backend::AbstractADType,
 )
-    (; conditions, backends, prep_Bᵀ) = implicit
+    (; conditions, backends) = implicit
+    (; prep_Bᵀ) = prep
     actual_backend = isnothing(backends) ? suggested_backend : backends.x
     contexts = (Constant(y), Constant(z), map(Constant, args)...)
     f_vec = VecToVec(conditions, x)
@@ -232,7 +249,8 @@ function build_Bᵀ(
             f_vec, prep_Bᵀ, actual_backend, x_vec, (dc_vec,), contexts...
         )
     end
-    function Bᵀ_fun(dc_vec)
+    function Bᵀ_fun(dc_vec_wrongtype)
+        copyto!(dc_vec, dc_vec_wrongtype)
         return pullback(f_vec, prep_Bᵀ_same, actual_backend, x_vec, (dc_vec,), contexts...)[1]
     end
     return Bᵀ_fun