# Unconstrained optimization

In this example, we show how to differentiate through the solution of the following unconstrained optimization problem:

$$$\hat{y}(x) = \min_{y \in \mathbb{R}^m} f(x, y)$$$

The optimality conditions are given by gradient stationarity:

$$$F(x, \hat{y}(x)) = 0 \quad \text{with} \quad F(x,y) = \nabla_2 f(x, y) = 0$$$
using ImplicitDifferentiation
using Optim
using Random
using Zygote

Random.seed!(63)
Random.TaskLocalRNG()

## Implicit function wrapper

To make verification easy, we minimize a quadratic objective

$$$f(x, y) = \lVert y - x \rVert^2$$$

In this case, the optimization algorithm is very simple (the identity function does the job), but still we implement it using a black box solver from Optim.jl to show that it doesn't change the result.

function dumb_identity(x)
f(y) = sum(abs2, y-x)
y0 = zero(x)
res = optimize(f, y0, LBFGS(); autodiff=:forward)
y = Optim.minimizer(res)
return y
end;

On the other hand, optimality conditions should be provided explicitly whenever possible, so as to avoid nesting autodiff calls.

zero_gradient(x, y) = 2(y - x);

We now have all the ingredients to construct our implicit function.

implicit = ImplicitFunction(dumb_identity, zero_gradient);

## Testing

x = rand(3, 2)
3×2 Matrix{Float64}:
0.224421  0.776206
0.326728  0.349684
0.621517  0.181591

Let's start by taking a look at the forward pass, which should be the identity function.

implicit(x)
3×2 Matrix{Float64}:
0.224421  0.776206
0.326728  0.349684
0.621517  0.181591

We now check whether the behavior of our ImplicitFunction wrapper is coherent with the theoretical derivatives.

Zygote.jacobian(implicit, x)
6×6 Matrix{Float64}:
1.0  -0.0  -0.0  -0.0  -0.0  -0.0
-0.0   1.0  -0.0  -0.0  -0.0  -0.0
-0.0  -0.0   1.0  -0.0  -0.0  -0.0
-0.0  -0.0  -0.0   1.0  -0.0  -0.0
-0.0  -0.0  -0.0  -0.0   1.0  -0.0
-0.0  -0.0  -0.0  -0.0  -0.0   1.0

As expected, we recover the identity matrix as Jacobian. Strictly speaking, the Jacobian should be a 4D tensor, but it is flattened by Zygote into a 2D matrix.

Note that implicit differentiation was necessary here, since our solver alone doesn't support autodiff with Zygote.jl.

try; Zygote.jacobian(dumb_identity, x); catch e; @error e; end
┌ Error: Zygote.CompileError(Tuple{typeof(Optim.perform_linesearch!), Optim.LBFGSState{Matrix{Float64}, Vector{Matrix{Float64}}, Vector{Matrix{Float64}}, Float64, Matrix{Float64}}, Optim.LBFGS{Nothing, LineSearches.InitialStatic{Float64}, LineSearches.HagerZhang{Float64, Base.RefValue{Bool}}, Optim.var"#19#21"}, Optim.ManifoldObjective{NLSolversBase.OnceDifferentiable{Float64, Matrix{Float64}, Matrix{Float64}}}}, ErrorException("try/catch is not supported."))
└ @ Main 1_unconstrained_optimization.md:82