Skip to content
New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Support Enzyme forward + reverse mode #10

Merged
merged 4 commits into from
Jul 25, 2023
Merged

Support Enzyme forward + reverse mode #10

merged 4 commits into from
Jul 25, 2023

Conversation

devmotion
Copy link
Member

@devmotion devmotion commented Jul 18, 2023

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).

@test adtype isa AutoForwardDiff
@test adtype isa AutoForwardDiff{nothing}

adtype = AutoForwardDiff(10)
Copy link
Member Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

In contrast to all other constructors, the convenience constructor for AutoForwardDiff uses positional instead of keyword arguments. Was this done on purpose?

Copy link
Member

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

I think not, it was carried over from previous definition in Optimization.jl as is. Feel free to fix it in this PR

@Vaibhavdixit02
Copy link
Member

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
Copy link
Member Author

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.

@Vaibhavdixit02 Vaibhavdixit02 merged commit 9e91d8a into main Jul 25, 2023
3 checks passed
@Vaibhavdixit02 Vaibhavdixit02 deleted the dw/enzyme branch July 25, 2023 13:53
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment
Labels
None yet
Projects
None yet
Development

Successfully merging this pull request may close these issues.

Enzyme types for both forward and reverse mode?
2 participants