Introduction

We explain the basics of our package on a simple function that is not amenable to naive automatic differentiation.

using ForwardDiff
using ImplicitDifferentiation
using LinearAlgebra
using Zygote

Why do we bother?

ForwardDiff.jl and Zygote.jl are two prominent packages for automatic differentiation in Julia. While they are very generic, there are simple language constructs that they cannot differentiate through.

function badsqrt(x::AbstractArray)
    a = [0.0]
    a[1] = x[1]
    return sqrt.(x)
end;

This is essentially the componentwise square root function but with an additional twist: a::Vector{Float64} is created internally, and its only element is replaced with the first element of x. We can check that it does what it's supposed to do.

x = [4.0, 9.0]
badsqrt(x)

2-element Vector{Float64}:
 2.0
 3.0

Of course the Jacobian has an explicit formula.

J = Diagonal(0.5 ./ sqrt.(x))

2×2 LinearAlgebra.Diagonal{Float64, Vector{Float64}}:
 0.25   ⋅ 
  ⋅    0.166667

However, things start to go wrong when we compute it with autodiff, due to the limitations of ForwardDiff.jl and those of Zygote.jl.

try
    ForwardDiff.jacobian(badsqrt, x)
catch e
    e
end

MethodError(Float64, (Dual{ForwardDiff.Tag{typeof(Main.badsqrt), Float64}}(4.0,1.0,0.0),), 0x0000000000007b47)

ForwardDiff.jl throws an error because it tries to call badsqrt with an array of dual numbers, and cannot use one of these numbers to fill a (which has element type Float64).

try
    Zygote.jacobian(badsqrt, x)
catch e
    e
end

ErrorException("Mutating arrays is not supported -- called setindex!(Vector{Float64}, ...)\nThis error occurs when you ask Zygote to differentiate operations that change\nthe elements of arrays in place (e.g. setting values with x .= ...)\n\nPossible fixes:\n- avoid mutating operations (preferred)\n- or read the documentation and solutions for this error\n  https://fluxml.ai/Zygote.jl/latest/limitations\n")

Zygote.jl also throws an error because it cannot handle mutation.

Implicit function

The first possible use of ImplicitDifferentiation.jl is to overcome the limitations of automatic differentiation packages by defining functions (and computing their derivatives) implicitly. An implicit function is a mapping

\[x \in \mathbb{R}^n \longmapsto y(x) \in \mathbb{R}^m\]

whose output is defined by conditions

\[c(x,y(x)) = 0 \in \mathbb{R}^m\]

We represent it using a type called ImplicitFunction, which you will see in action shortly.

First we define a forward mapping corresponding to the function we consider. It returns the actual output $y(x)$ of the function, and can be thought of as a black box solver. Importantly, this Julia callable doesn't need to be differentiable by automatic differentiation packages but the underlying function still needs to be mathematically differentiable.

forward(x) = badsqrt(x);

Then we define conditions $c(x, y) = 0$ that the output $y(x)$ is supposed to satisfy. These conditions must be array-valued, with the same size as $y$. Unlike the forward mapping, the conditions need to be differentiable by automatic differentiation packages with respect to both $x$ and $y$. Here the conditions are very obvious: the square of the square root should be equal to the original value.

function conditions(x, y)
    c = y .^ 2 .- x
    return c
end;

Finally, we construct a wrapper implicit around the previous objects. By default, forward is assumed to return a single output and conditions is assumed to accept 2 arguments.

implicit = ImplicitFunction(forward, conditions)

ImplicitFunction{lazy}(forward, conditions, ImplicitDifferentiation.KrylovLinearSolver(), nothing, nothing)

What does this wrapper do? When we call it as a function, it just falls back on implicit.forward, so unsurprisingly we get the output $y(x)$.

implicit(x)

2-element Vector{Float64}:
 2.0
 3.0

And when we try to compute its Jacobian, the implicit function theorem is applied in the background to circumvent the lack of differentiability of the forward mapping.

Forward and reverse mode autodiff

Now ForwardDiff.jl works seamlessly.

ForwardDiff.jacobian(implicit, x) ≈ J

true

And so does Zygote.jl. Hurray!

Zygote.jacobian(implicit, x)[1] ≈ J

true

This page was generated using Literate.jl.