Tricks

We demonstrate several features that may come in handy for some users.

using ComponentArrays
using ForwardDiff
using ImplicitDifferentiation
using Krylov
using LinearAlgebra
using Zygote

ComponentArrays

For when you need derivatives with respect to multiple inputs or outputs.

function forward_components_aux(a::AbstractVector, b::AbstractVector, m::Number)
    d = m * sqrt.(a)
    e = sqrt.(b)
    return d, e
end

function conditions_components_aux(a, b, m, d, e)
    c_d = (d ./ m) .^ 2 .- a
    c_e = (e .^ 2) .- b
    return c_d, c_e
end;

You can use ComponentVector from ComponentArrays.jl as an intermediate storage.

function forward_components(x::ComponentVector)
    d, e = forward_components_aux(x.a, x.b, x.m)
    y = ComponentVector(; d=d, e=e)
    return y
end

function conditions_components(x::ComponentVector, y::ComponentVector)
    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)
    return c
end;

And build your implicit function like so.

implicit_components = ImplicitFunction(forward_components, conditions_components);

Since ComponentVectors are not yet compatible with iterative solvers from Krylov.jl, we (temporarily) need a bit of type piracy to make it work

Krylov.ktypeof(::ComponentVector{T,V}) where {T,V} = V

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)

ComponentVector{Float64}(d = [6.0, 8.485281374238571], e = [1.7320508075688772, 2.0, 2.23606797749979])

And it works with both ForwardDiff.jl and Zygote.jl

ForwardDiff.jacobian(implicit_components, x)

5×6 Matrix{Float64}:
 3.0  0.0      0.0       0.0   0.0       1.0
 0.0  2.12132  0.0       0.0   0.0       1.41421
 0.0  0.0      0.288675  0.0   0.0       0.0
 0.0  0.0      0.0       0.25  0.0       0.0
 0.0  0.0      0.0       0.0   0.223607  0.0

Zygote.jacobian(implicit_components, x)[1]

5×6 Matrix{Float64}:
 3.0  0.0      0.0       0.0   0.0       1.0
 0.0  2.12132  0.0       0.0   0.0       1.41421
 0.0  0.0      0.288675  0.0   0.0       0.0
 0.0  0.0      0.0       0.25  0.0       0.0
 0.0  0.0      0.0       0.0   0.223607  0.0

function full_pipeline(a, b, m)
    x = ComponentVector(; a=a, b=b, m=m)
    y = implicit_components(x)
    return y.d, y.e
end;

Byproducts

For when you need an additional output but don't care about its derivative.

function forward_byproduct(x)
    z = 2  # "randomized" choice
    y = x .^ (1 / z)
    return y, z
end

function conditions_byproduct(x, y, z)
    c = y .^ z .- x
    return c
end;

The syntax for building the implicit function is the same.

implicit_byproduct = ImplicitFunction(forward_byproduct, conditions_byproduct);

But this time the return value is a tuple (y, z)

x = [4.0, 9.0]
implicit_byproduct(x)

([2.0, 3.0], 2)

And it works with both ForwardDiff.jl and Zygote.jl

ForwardDiff.jacobian(first ∘ implicit_byproduct, x)

2×2 Matrix{Float64}:
 0.25  0.0
 0.0   0.166667

Zygote.jacobian(first ∘ implicit_byproduct, x)[1]

2×2 Matrix{Float64}:
  0.25  -0.0
 -0.0    0.166667

This page was generated using Literate.jl.