Replace internal AD backend types with ADTypes #2047

devmotion · 2023-07-19T00:57:43Z

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

torfjelde · 2023-07-19T09:52:12Z

Love it!

torfjelde · 2023-07-19T09:54:24Z

We'll also need to make equivalent changes in AdvancedVI, I believe.

yebai · 2023-07-19T10:56:59Z

We'll also need to make equivalent changes in AdvancedVI, I believe.

I think @Red-Portal already did it in the AdvancedVI rewrite PR.

Red-Portal · 2023-07-19T11:06:30Z

Hi, yes that's already up to date!

docs/src/using-turing/autodiff.md

TuringLang/Turing.jl#2047

src/essential/ad.jl

Project.toml

src/essential/Essential.jl

Project.toml

devmotion · 2023-11-15T08:38:16Z

I opened tpapp/LogDensityProblemsAD.jl#21.

devmotion · 2023-11-15T08:41:11Z

src/essential/Essential.jl

+        AutoForwardDiff,
+        AutoTracker,
+        AutoZygote,
+        AutoReverseDiff,


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

src/essential/ad.jl

sunxd3 · 2023-11-15T09:33:03Z

test/essential/ad.jl

-        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


@devmotion I had to pass in θ for the result to be correct, started a PR tpapp/LogDensityProblemsAD.jl#22

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.

Project.toml

yebai · 2023-11-15T17:02:48Z

src/mcmc/hmc.jl

 ) where AD
-    return HMC{AD}(ϵ, n_leapfrog, metricT, space)
+    return HMC(ϵ, n_leapfrog, metricT, space; adtype = adtype)


@devmotion, we can probably consider removing the global AD flag ADBACKEND, and always specify the autodiff backend in inference algorithms. What are your thoughts?

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?

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.

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.

github-actions · 2023-11-16T09:33:35Z

Pull Request Test Coverage Report for Build 6887921533

0 of 41 (0.0%) changed or added relevant lines in 5 files are covered.
1 unchanged line in 1 file lost coverage.
Overall coverage remained the same at 0.0%

Changes Missing Coverage	Changed/Added Lines	%
src/mcmc/Inference.jl	1	0.0%
ext/TuringDynamicHMCExt.jl	2	0.0%
src/mcmc/sghmc.jl	4	0.0%
src/essential/ad.jl	12	0.0%
src/mcmc/hmc.jl	22	0.0%

Files with Coverage Reduction	New Missed Lines	%
ext/TuringDynamicHMCExt.jl	1	0.0%

Totals
Change from base Build 6829598411:	0.0%
Covered Lines:	0
Relevant Lines:	1421

💛 - Coveralls

codecov · 2023-11-16T09:34:09Z

Codecov Report

Attention: 41 lines in your changes are missing coverage. Please review.

Comparison is base (d4a7975) 0.00% compared to head (da44611) 0.00%.

Files	Patch %	Lines
src/mcmc/hmc.jl	0.00%	22 Missing ⚠️
src/essential/ad.jl	0.00%	12 Missing ⚠️
src/mcmc/sghmc.jl	0.00%	4 Missing ⚠️
ext/TuringDynamicHMCExt.jl	0.00%	2 Missing ⚠️
src/mcmc/Inference.jl	0.00%	1 Missing ⚠️

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.
📢 Have feedback on the report? Share it here.

devmotion · 2023-11-16T09:37:28Z

src/mcmc/hmc.jl

-    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}


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

We plan to depreciate and then remove this old interface once the AbstractMCMC-based externalsampler interface works for Gibbs.

yebai · 2023-11-16T11:38:52Z

CI errors about chain resume/save are unrelated to this PR.

yebai

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 <[email protected]> * 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 <[email protected]>

storopoli · 2023-12-20T15:45:43Z

So the tape gets compiled by default now when using ReverseDiff?
No need more for Turing.setrdcache(true)?

yebai · 2023-12-20T15:52:28Z

You need to pass AutoReverseDiff(true) (true enables compiled tape, false disables it) to individual sampling algorithms, see more at https://github.com/TuringLang/Turing.jl/blob/master/HISTORY.md

storopoli · 2023-12-20T16:37:33Z

Yes, but still is not clear. The example has only adtype=AutoForwardDiff(; chunksize)) in the sampler constructor.
I had no idea that AutoReverseDiff took positional arguments, and why AutoForwardDiff doesn't take any positional arguments?

Also, https://github.com/SciML/ADTypes.jl doesn't have docs, so it is even harder to figure it out.

yebai · 2023-12-20T16:42:25Z

@sunxd3, maybe add a few concrete examples for each popular autodiff backend to HISTORY.md and docstrings?

sunxd3 · 2023-12-20T16:44:21Z

Yeah, it is confusing, I'll have a PR

storopoli · 2023-12-20T16:49:33Z

I can do that if you want to. I am hours away from my Holiday break.
At least the docs part.

EDIT: here (https://turinglang.org/v0.30/docs/using-turing/autodiff)
EDIT 2: already done here (TuringLang/docs#430) huh? Oh that's great then.

sunxd3 · 2023-12-20T17:09:59Z

@storopoli yep, the tutorial is updated. I just updated the release information too

Replace internal AD backend types with ADTypes

91f9715

yebai reviewed Jul 30, 2023

View reviewed changes

docs/src/using-turing/autodiff.md Outdated Show resolved Hide resolved

yebai force-pushed the master branch from 8f61144 to c2c2a70 Compare July 31, 2023 16:00

Red-Portal mentioned this pull request Aug 26, 2023

Update ADVI interface to match AdvancedVI.jl #2068

Closed

sunxd3 added a commit to TuringLang/docs that referenced this pull request Nov 14, 2023

Update autodiff.jmd following adaptation of ADTypes

e03f5d3

TuringLang/Turing.jl#2047

sunxd3 mentioned this pull request Nov 14, 2023

Update autodiff.jmd following adaptation of ADTypes TuringLang/docs#430

Merged

Merge remote-tracking branch 'origin/master' into dw/adtypes

347c879

devmotion commented Nov 14, 2023

View reviewed changes

src/essential/ad.jl Outdated Show resolved Hide resolved

devmotion commented Nov 14, 2023

View reviewed changes

src/essential/ad.jl Outdated Show resolved Hide resolved

devmotion commented Nov 14, 2023

View reviewed changes

Project.toml Outdated Show resolved Hide resolved

devmotion commented Nov 14, 2023

View reviewed changes

src/essential/Essential.jl Outdated Show resolved Hide resolved

devmotion and others added 3 commits November 14, 2023 14:56

Remove upstreamed functionality

d01b9ee

Update ADBackend code

ee19068

Formatting

8a956c4

devmotion commented Nov 15, 2023

View reviewed changes

Project.toml Outdated Show resolved Hide resolved

Update Project.toml

773e826

devmotion commented Nov 15, 2023

View reviewed changes

src/essential/ad.jl Outdated Show resolved Hide resolved

devmotion and others added 3 commits November 15, 2023 09:50

Update src/essential/ad.jl

e711c86

Fix tests

4a58f69

Switch ADType version to 0.1.5 for CI testing

834ba48

sunxd3 reviewed Nov 15, 2023

View reviewed changes

A few fixes

7398d96

devmotion commented Nov 15, 2023

View reviewed changes

Project.toml Outdated Show resolved Hide resolved

Update Project.toml

a7a2d23

devmotion added 2 commits November 15, 2023 14:28

Make ad type a field

cf73afd

Improve docs and tests

b8cf60e

yebai reviewed Nov 15, 2023

View reviewed changes

sunxd3 added 2 commits November 16, 2023 05:17

small fix

1b9934c

Fix test errors

da44611

devmotion commented Nov 16, 2023

View reviewed changes

yebai marked this pull request as ready for review November 16, 2023 11:39

yebai approved these changes Nov 16, 2023

View reviewed changes

yebai merged commit 6649f10 into master Nov 16, 2023
11 of 13 checks passed

yebai deleted the dw/adtypes branch November 16, 2023 12:23

yebai mentioned this pull request Nov 16, 2023

Remove global AD flag ADBACKEND #2132

Closed

sunxd3 mentioned this pull request Nov 24, 2023

Remove global AD flag "ADBACKEND" and function suite #2134

Merged

4 tasks

ya0 mentioned this pull request Jan 18, 2024

Docs: setadbackend #2155

Closed

Provide feedback

Saved searches

Use saved searches to filter your results more quickly

Replace internal AD backend types with ADTypes #2047

Replace internal AD backend types with ADTypes #2047

devmotion commented Jul 19, 2023

torfjelde commented Jul 19, 2023

torfjelde commented Jul 19, 2023

yebai commented Jul 19, 2023

Red-Portal commented Jul 19, 2023

devmotion commented Nov 15, 2023

devmotion Nov 15, 2023

sunxd3 Nov 15, 2023

devmotion Nov 15, 2023

yebai Nov 15, 2023

devmotion Nov 16, 2023

yebai Nov 16, 2023

devmotion Nov 16, 2023

github-actions bot commented Nov 16, 2023

codecov bot commented Nov 16, 2023

devmotion Nov 16, 2023

yebai Nov 16, 2023 •

edited

Loading

yebai commented Nov 16, 2023

yebai left a comment

storopoli commented Dec 20, 2023

yebai commented Dec 20, 2023

storopoli commented Dec 20, 2023

yebai commented Dec 20, 2023

sunxd3 commented Dec 20, 2023

storopoli commented Dec 20, 2023 •

edited

Loading

sunxd3 commented Dec 20, 2023

Replace internal AD backend types with ADTypes #2047

Replace internal AD backend types with ADTypes #2047

Conversation

devmotion commented Jul 19, 2023

torfjelde commented Jul 19, 2023

torfjelde commented Jul 19, 2023

yebai commented Jul 19, 2023

Red-Portal commented Jul 19, 2023

devmotion commented Nov 15, 2023

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

github-actions bot commented Nov 16, 2023

Pull Request Test Coverage Report for Build 6887921533

💛 - Coveralls

codecov bot commented Nov 16, 2023

Codecov Report

Choose a reason for hiding this comment

yebai Nov 16, 2023 • edited Loading

Choose a reason for hiding this comment

yebai commented Nov 16, 2023

yebai left a comment

Choose a reason for hiding this comment

storopoli commented Dec 20, 2023

yebai commented Dec 20, 2023

storopoli commented Dec 20, 2023

yebai commented Dec 20, 2023

sunxd3 commented Dec 20, 2023

storopoli commented Dec 20, 2023 • edited Loading

sunxd3 commented Dec 20, 2023

yebai Nov 16, 2023 •

edited

Loading

storopoli commented Dec 20, 2023 •

edited

Loading