How to write a program differentiably
In this blog, I will show how to write a linear algebra function, the sparse matrix-vector multiplication function, differentiably, by converting an existing irreversible program in SparseArrays to a reversible one, step by step.
First of all, you need to install the package of reversible embedded domain-specific language (eDSL) NiLang (v0.9), by typing
Step1: Find the function that you want to differentiate
As an example, I will re-write the sparse matrix-vector multiplication function defined in
The following is a commented version of this program. I put some remarks above the "evil" statements that make a program irreversible so that we can fix them later.
Step2: Modify the irreversible statements
Finally, your reversible program will look like
Step3: Test and benchmark your program
A reversible program written in this way is also differentiable. We can benchmark the automatic differentiation by typing
Here, we used the
GVar type to store the gradients, which is similar to the
Dual type in ForwardDiff.jl, but in reverse mode. The above benchmark shows the performance of back propagating the gradients. It is easy to see complex-valued autodiff is supported in NiLang. In the practical case, one may want to define a real-valued output as the loss functions. We recommend readers to read our paper for more detials.
NiLang is still experimental, with a lot of unexpected "features". Feel free to join and discuss in Julia slack with me. I appear in the #autodiff and #reversible-computing channel quite often.