EnzymeAD / Enzyme.jl

Julia bindings for the Enzyme automatic differentiator
https://enzyme.mit.edu
MIT License
456 stars 64 forks source link

Utilities to translate from `ReverseMode` to `ReverseModeSplit` #1938

Open gdalle opened 1 month ago

gdalle commented 1 month ago

In DI, the pullback functions use split reverse mode by default, so they disregard the mode field of AutoEnzyme (because I figure no one is going to pass a ReverseSplitWithPrimal anyway). But still I'd like to respect some of the attributes of the ReverseMode that the user provides. Would you welcome a PR with functions that translate a ReverseMode to the corresponding ReverseModeSplit, setting reasonable defaults for the remaining type parameters?

wsmoses commented 1 month ago

Sure!

wsmoses commented 1 month ago

Though you should still check the mode in case they do

gdalle commented 1 month ago

Then I'd also need translation in the other direction, in case they pass me a split mode but I just wanna do gradient

wsmoses commented 1 month ago

Also sure if you want