Replace internal AD backend types with ADTypes#2047
Conversation
|
Love it! |
|
We'll also need to make equivalent changes in AdvancedVI, I believe. |
I think @Red-Portal already did it in the AdvancedVI rewrite PR. |
|
Hi, yes that's already up to date! |
|
I opened tpapp/LogDensityProblemsAD.jl#21. |
| AutoForwardDiff, | ||
| AutoTracker, | ||
| AutoZygote, | ||
| AutoReverseDiff, |
There was a problem hiding this comment.
Do we actually want to export these types? Or should we tell users to use ADTypes.AutoForwardDiff etc. (in particular when these types would be used in other packages such as e.g. AdvancedVI as well).
| f_rd = LogDensityProblemsAD.ADgradient(Turing.Essential.ReverseDiffAD{false}(), f) | ||
| f_rd_compiled = LogDensityProblemsAD.ADgradient(Turing.Essential.ReverseDiffAD{true}(), f) | ||
| f_rd = LogDensityProblemsAD.ADgradient(Turing.Essential.ReverseDiffAD(false), f) | ||
| f_rd_compiled = LogDensityProblemsAD.ADgradient(Val(:ReverseDiff), f; compile=Val(true), x=θ) # need to compile with non-zero inputs |
There was a problem hiding this comment.
@devmotion I had to pass in θ for the result to be correct, started a PR tpapp/LogDensityProblemsAD.jl#22
There was a problem hiding this comment.
I think we should not add a keyword argument but just continue overloading ADgradient for TuringLogDensityFunction. I wouldn't want users to have to deal with such a keyword argument.
| adtype::ADTypes.AbstractADType = ADBackend(), | ||
| ) where AD | ||
| return HMC{AD}(ϵ, n_leapfrog, metricT, space) | ||
| return HMC(ϵ, n_leapfrog, metricT, space; adtype = adtype) |
There was a problem hiding this comment.
@devmotion, we can probably consider removing the global AD flag ADBACKEND, and always specify the autodiff backend in inference algorithms. What are your thoughts?
There was a problem hiding this comment.
And set the default AD type for each algorithm/sampler to a specific default such as AutoForwardDiff, you mean? Or would you like users to always specify the AD type explicitly?
There was a problem hiding this comment.
Yes, we can set the default AD type to AutoForwardDiff for all algorithms and allow users to override them via keyword arguments. That way, we don't need to maintain a global AD flag and can remove the messy code around it. But that should probably be done in a separate PR.
There was a problem hiding this comment.
I would like to get rid of the global flag 👍 If we make it in a separate PR we should maybe wait with tagging a breaking release until this follow-up PR is merged as well, to avoid two breaking releases in a row.
Pull Request Test Coverage Report for Build 6887921533
💛 - Coveralls |
Codecov ReportAttention:
Additional details and impacted files@@ Coverage Diff @@
## master #2047 +/- ##
======================================
Coverage 0.00% 0.00%
======================================
Files 21 21
Lines 1435 1421 -14
======================================
+ Misses 1435 1421 -14 ☔ View full report in Codecov by Sentry. |
| kwargs... | ||
| ) where AD | ||
| return HMC{AD}(ϵ, n_leapfrog; kwargs...) | ||
| function HMC(ϵ::Float64, n_leapfrog::Int, ::Type{metricT}, space::Tuple; adtype::ADTypes.AbstractADType = ADBackend()) where {metricT <: AHMC.AbstractMetric} |
There was a problem hiding this comment.
Should we remove this constructor and only support metricT as keyword argument? Or make all arguments keyword arguments?
Generally, these HMC constructors are quite messy...
There was a problem hiding this comment.
We plan to depreciate and then remove this old interface once the AbstractMCMC-based externalsampler interface works for Gibbs.
|
CI errors about |
yebai
left a comment
There was a problem hiding this comment.
Thanks @sunxd3 @devmotion -- good work!
* Update autodiff.jmd following adaptation of `ADTypes` TuringLang/Turing.jl#2047 * Update tutorials/docs-10-using-turing-autodiff/autodiff.jmd Co-authored-by: github-actions[bot] <41898282+github-actions[bot]@users.noreply.github.com> * Using `ADTypes` for ad doc * Update tutorials/docs-10-using-turing-autodiff/autodiff.jmd Co-authored-by: github-actions[bot] <41898282+github-actions[bot]@users.noreply.github.com> * Update tutorials/docs-10-using-turing-autodiff/autodiff.jmd Co-authored-by: github-actions[bot] <41898282+github-actions[bot]@users.noreply.github.com> * Apply suggestions from code review Co-authored-by: Hong Ge <3279477+yebai@users.noreply.github.com> * Update tutorials/docs-10-using-turing-autodiff/autodiff.jmd * Update tutorials/docs-10-using-turing-autodiff/autodiff.jmd Co-authored-by: github-actions[bot] <41898282+github-actions[bot]@users.noreply.github.com> * Apply suggestions from code review * Update tutorials/docs-10-using-turing-autodiff/autodiff.jmd * Update tutorials/docs-10-using-turing-autodiff/autodiff.jmd * Update tutorials/docs-10-using-turing-autodiff/autodiff.jmd * Update tutorials/docs-10-using-turing-autodiff/autodiff.jmd Co-authored-by: github-actions[bot] <41898282+github-actions[bot]@users.noreply.github.com> --------- Co-authored-by: github-actions[bot] <41898282+github-actions[bot]@users.noreply.github.com> Co-authored-by: Hong Ge <3279477+yebai@users.noreply.github.com>
|
So the tape gets compiled by default now when using |
|
You need to pass |
|
Yes, but still is not clear. The example has only Also, https://github.com/SciML/ADTypes.jl doesn't have docs, so it is even harder to figure it out. |
|
@sunxd3, maybe add a few concrete examples for each popular autodiff backend to |
|
Yeah, it is confusing, I'll have a PR |
|
I can do that if you want to. I am hours away from my Holiday break. EDIT: here (https://turinglang.org/v0.30/docs/using-turing/autodiff) |
|
@storopoli yep, the tutorial is updated. I just updated the release information too |
This PR is a draft proposal for replacing our internal AD backend types with ADTypes.
Needs support of ADTypes in LogDensityProblemsAD: tpapp/LogDensityProblemsAD.jl#17