SciML / ADTypes.jl

Repository for automatic differentiation backend types
https://sciml.github.io/ADTypes.jl/
MIT License
37 stars 11 forks source link

Support Enzyme forward + reverse mode #10

Closed devmotion closed 1 year ago

devmotion commented 1 year ago

Fixes #9 and adds more tests.

I chose to not introduce any new types but instead added a field to AutoEnzyme (and hence an additional type parameter which makes the PR technically breaking but I'm not sure how widespread the use of AutoEnzyme is within the community and if there are any examples that would actually break?). The advantage is that one could even specify other modes than Enzyme.Forward and Enzyme.Reverse such as Enzyme.ReverseWithPrimal.

Since I set the default value to EnzymeCore.Reverse, a dependency on EnzymeCore is added. Even though it is lightweight (only depends on Adapt and I'm about to open a PR that would make it a weak dependency on Julia >= 1.9, so if accepted the package would have zero dependencies on Julia >= 1.9) this might not be desirable. An alternative would be to use a default value of nothing instead and to not explicitly restrict the field to EnzymeCore.Mode.

Edit: In CI I noticed another issue with EnzymeCore - it breaks compatibility with Julia < 1.6. On the other hand, in these versions AutoEnzyme is not useful at all since Enzyme >= 0.4 (currently at 0.11.5) does only support Julia >= 1.6.

Edit 2: Yet another possibility would be to not restrict the type of the field and move the convenience constructor to an extension such that at least on Julia >= 1.9 no hard dependencies are added while still keeping a more meaningful default than nothing. On Julia < 1.9 and >= 1.6, this constructor could be hidden behind a requires-block or EnzymeCore could be kept as a direct dependency in these older versions. On Julia < 1.6, the convenience constructor could still not be supported (but as mentioned on these the whole AutoEnzyme type is not useful anyway).

Vaibhavdixit02 commented 1 year ago

Yet another possibility would be to not restrict the type of the field and move the convenience constructor to an extension such that at least on Julia >= 1.9 no hard dependencies are added while still keeping a more meaningful default than nothing. On Julia < 1.9 and >= 1.6, this constructor could be hidden behind a requires-block or EnzymeCore could be kept as a direct dependency in these older versions. On Julia < 1.6, the convenience constructor could still not be supported (but as mentioned on these the whole AutoEnzyme type is not useful anyway).

I like this the best. We can retain the parametric type you have added but remove the EnzymeCore.Mode subtyping for it.

devmotion commented 1 year ago

On a second thought, I think maybe a default of nothing is even more desirable. This will allow downstream packages to pick a desired default themselves and is more consistent with other AD types such as AutoFiniteDifferences.