Jacobian row partitioning and reverse-mode AD

[ChrisRackauckas]

Reverse-mode AD implementations calculate a Jacobian row by row instead of column by column. Thus it would be nice to have a way to do row-wise matrix partitioning and coloring. Once we have these color vectors, we have to specialize the Jacobian function in Zygote to allow for and exploit a color vector.

[pkj-m]

Are we supposed to write the function for reverse mode AD from scratch like we did for forward mode or modify some pre-existing files in Zygote?

[ChrisRackauckas]

Let's at least prototype in this repo, and then think about upstreaming it.

[ChrisRackauckas]

You can use https://github.com/FluxML/Tracker.jl/blob/master/src/back.jl#L169-L178 as a template.

[gdalle]

Apologies in advance for the repeated notifications, I'm going through every SparseDiffTools issue to see which ones are solved by SparseConnectivityTracer + SparseMatrixColorings + DifferentiationInterface.

This one is supported out of the box: if the user chooses a reverse mode backend, the coloring and compressed differentiation will automatically happen row-wise.

https://gdalle.github.io/SparseMatrixColorings.jl/stable/api/#ADTypes.row_coloring