Navigating the automatic differentiation ecosystem in Julia

Question

Julia has a somewhat sprawling AD ecosystem, with perhaps by now more than a dozen different packages spanning, as far as I can tell, forward-mode (ForwardDiff.jl, ForwardDiff2.jl ), reverse-mode (ReverseDiff.jl, Nabla.jl, AutoGrad.jl), and source-to-source (Zygote.jl, Yota.jl, Enzyme.jl, and presumably also the forthcoming Diffractor.jl) at several different steps of the compilation pipeline, as well as more exotic things like NiLang.jl.

Between such packages, what is the support for different language constructs (control-flow, mutation, etc.), and are there any rules of thumb for how one should go about choosing a given AD for a given task? I believe there was a compare-and-contrast table on the Julia Slack at some point, but I can't seem find anything like that reproduced for posterity in the relevant discourse threads or other likely places (1, 2)

Bailiff

Answer 1

I'd also love to hear an informed answer to this. Some more links that might be of interest.

Diffractor now has a Github repo, which lays out the implementation plan. After reading the text there, my take is that it will require long-term implementation work before Diffractor is production ready. On the other hand, there is a feeling that Zygote may be in "maintenance mode" while awaiting Diffractor. At least from a distance, the situation seems a bit awkward. The good news is that the ChainRules.jl ecosystem seems to make it possible to easily swap between autodiff systems.

As of Sept 2021, Yota seems to be rapidly evolving. The 0.5 release brings support for ChainRules which seems to unlock it for production use. There is a lot of interesting discussion at this release thread. My understanding from reading through those threads is that the scope of Yota is more limited compared to Zygote (e.g., autodiff through mutation is not supported). This limited scope has the advantage of opening up optimization opportunities such as preallocation, and kernel fusion that may not be possible in a more general autodiff system. As such, Yota might be better suited to fill the niche of, e.g., PyTorch type modeling.

Rapacious