tpapp / LogDensityProblemsAD.jl

AD backends for LogDensityProblems.jl.
MIT License
12 stars 6 forks source link

Fix Enzyme implementation for 0.13 #38

Closed ptiede closed 1 month ago

ptiede commented 1 month ago

The Enzyme Ext for Enzyme version 0.13 is currently incorrect. A number of things changed that makes the current version on main incorrect including

In addition if you need runtime activity you now need to pass <:ReverseMode{true}. The EnzymeExt on main will drop this and always go with ReverseMode{false}.

This PR tries to fix all of these issues by

  1. Explicitly passing the ReverseMode and ForwardMode into ADgradient and converting them to requiring the primal if needed
  2. Fixing the ordering for ForwardMode in logdensity_and_gradient
ptiede commented 1 month ago

After talking with @wsmoses I'm going to move some of the set_primal stuff into EnzymeCore

tpapp commented 1 month ago

@ptiede: thanks! so should we wait for some other change, or merge if tests pass?

FWIW, currently 1.9 can't satisfy version requirements, but I am happy to drop Julia 1.9 if @devmotion agrees. Users of 1.9 can still install earlier versions of the package.

ptiede commented 1 month ago

@devmotion I've updated the PR to utilize Enzyme's new WithPrimal feature. Let me know if you would prefer to continue to support 1.9. If so, I can the branches.

tpapp commented 1 month ago

FWIW, I would go for the simple solution and drop 1.9.

ptiede commented 1 month ago

No problem! Let me know if you need anything else!

tpapp commented 1 month ago

@devmotion: OK if we merge this?