Adjoint Sensitivities in Julia with Zygote & ChainRules

AD in Julia is fantastic. It works on almost all parts of the language and is highly efficient. Still, there are some computer codes that are hard to differentiate through. Commonly, solution algorithms that rely on iterative approximations (like root-finding methods for nonlinear (systems of) equations) can be hard to take gradients over. Firstly, this is caused by their iterative nature, which often results in AD packages unrolling all loop iterations and attempting to differentiate each iteration instead of moving over the entire loop. Secondly, some small sub function in the forward pass might be non-differentiable. That can either be because it is non-differentiable by itself (like calling external C code) or there is just no adjoint rule implemented for it yet.

-------

-------

Timestamps:
00:00 Intro
00:40 The (scalar) residual
01:12 A function to obtain the root
03:03 Implementing a loss functional
04:11 Attempt: Using Zygote but failing
06:45 Defining an Adjoint Rule to make it work
15:57 Restarting the Julia REPL session
16:42 Zygote gradient with the custom adjoint rule
17:20 Check derivative with Finite Differences
18:37 Summary
19:35 Outro