Autodiff backends

[lostella]

Right now the package falls back to Zygote for computing gradients using automatic differentiation. It would be nice to add support for multiple AD backends (ReverseDiff, Enzyme, Yota, Nabla, Diffractor… anything that works) so that one can chose what should be used.

This would require a mechanism for “optional dependencies”, so that a certain backend can be used only if the required dependency is installed.

Two ways I see for selecting what backend to use:

1. Some global variable
2. Wrapping functions in a specific type that enforces using this or that AD backend, eg, ZygoteFunction(f), ReverseDiffFunction(f), …

[nantonel]

isn't ChainRulesCore the AD core package?

[lostella]

The way I understand the current AD ecosystem (without looking much into the packages though): * ChainRulesCore defines the API and core utils for differentiation rules. * ChainRules defines differentiation rules for the stuff in Julia Base (I think?) based on the ChainRulesCore API. * AD “systems” like Zygote or others, implement the actual logic to compute derivatives; they may or may not be based on rules from the ChainRulesCore API, and if they do then it will be easy for people to add derivation rules for custom stuff. In any case, these are the packages one would use to actually compute derivatives. If someone from the Julia AD community happens to read this, feel free to confirm or correct me :D

…

Am 26.02.2022 um 01:08 schrieb Niccolò Antonello ***@***.***>:  isn't ChainRulesCore the AD core package? — Reply to this email directly, view it on GitHub, or unsubscribe. You are receiving this because you authored the thread.

[mohamed82008]

Check https://github.com/JuliaDiff/AbstractDifferentiation.jl and the associated paper https://arxiv.org/abs/2109.12449.

[lostella]

@mohamed82008 thanks, it seemed to much of a general issue not to have a solution already :-) now I’m curious, I’ll look into that

[mohamed82008]

Feel free to ping me if you have any questions :) I always wanted to add nonconvex optimisation support to ProximalAlgorithms but didn't get a chance to. If I understand correctly, nonconvex functions and constraints are now supported via AD?

[lostella]

@mohamed82008 ProximalAlgorithms supports nonconvex problems, in the sense that some of the implemented algorithms don’t need convexity assumptions on the objective terms in order to converge.

AD (Zygote right now) is used by default to compute gradients when they are needed. Constraints are handled via their projection mapping, and more generally non-differentiable terms are handled via proximal mappings (projections are a special case of this).

Here is a little bit of documentation on the primitives which are used, and that one can implement on custom types.

[mohamed82008]

Fantastic! Exactly what I was hoping for. Thanks!