Allow ADTypes AD backend selection in Hamiltonian #405

ErikQQY · 2025-03-16T17:56:14Z

While LogDensityProblemsAD.jl supports the AD backend from ADTypes.jl, we can also allow this for an unifying interface

ext/AdvancedHMCADTypesExt.jl

Project.toml

test/demo.jl

ext/AdvancedHMCADTypesExt.jl

test/demo.jl

gdalle · 2025-03-22T07:38:26Z

src/AdvancedHMC.jl

-        LogDensityProblemsAD.ADgradient(Val(:ForwardDiff), ℓ; kwargs...)
+        _logdensity = Base.Fix1(LogDensityProblems.logdensity, ℓ)
+        _logdensity_and_gradient = function (x)
+            prep = DI.prepare_gradient(_logdensity, adtype, x)


Preparing outside the closure would be more efficient, but users have to be warned I guess

I agree we should prepare outside the closure, but since we need the log density and its gradient of the model as a function in Hamiltonian, how can we prepare in advance?

Do you have access to a "typical input" somewhere?

AFAIK, the Hamiltonian is basically used for keeping the log density and its gradient function, and "typical input" is used in the subsequent sampling process which may not be possible to pass in here.

gdalle · 2025-03-22T07:40:13Z

src/AdvancedHMC.jl

@@ -146,7 +152,12 @@ function Hamiltonian(metric::AbstractMetric, ℓ; kwargs...)
    ℓπ = if cap === LogDensityProblems.LogDensityOrder{0}()
        # In this case ℓ does not support evaluation of the gradient of the log density function
        # We use ForwardDiff to compute the gradient
-        LogDensityProblemsAD.ADgradient(Val(:ForwardDiff), ℓ; kwargs...)
+        _logdensity = Base.Fix1(LogDensityProblems.logdensity, ℓ)


One can also use DI.Constant for ℓ if it contains no differentiable storage

gdalle · 2025-03-22T07:40:45Z

src/AdvancedHMC.jl

-function Hamiltonian(
-    metric::AbstractMetric, ℓπ::LogDensityModel, kind::Union{Symbol,Val,Module}; kwargs...
-)


I guess removing this is a breaking change?

test/contrib.jl

test/demo.jl

devmotion · 2025-03-22T21:54:16Z

Project.toml

 ProgressMeter = "92933f4c-e287-5a05-a399-4b506db050ca"
 Random = "9a3f8284-a2c9-5f02-9a11-845980a1fd5c"
+Reexport = "189a3867-3050-52da-a836-e630ba90ab69"


Unconditional reexports are a code smell IMO - why can't users just load ADTypes? By unconditional reexporting everything you completely lose control over what's defined in your package and automatically the whole API of the reexported package (and hence all breaking changes etc.) become part of the API of your package.

Yeah, users who want to try other AD backends should just load ADTypes by themselves then

devmotion · 2025-03-22T21:57:10Z

src/AdvancedHMC.jl

 end
-function Hamiltonian(metric::AbstractMetric, ℓ; kwargs...)
+function Hamiltonian(metric::AbstractMetric, ℓ; adtype=AutoForwardDiff(), kwargs...)


Suggested change

function Hamiltonian(metric::AbstractMetric, ℓ; adtype=AutoForwardDiff(), kwargs...)

function Hamiltonian(metric::AbstractMetric, ℓ; adtype::ADType=AutoForwardDiff(), kwargs...)

Suggested change

function Hamiltonian(metric::AbstractMetric, ℓ; adtype=AutoForwardDiff(), kwargs...)

function Hamiltonian(metric::AbstractMetric, ℓ; adtype=ADTypes.AutoForwardDiff(), kwargs...)

devmotion · 2025-03-22T21:57:19Z

src/AdvancedHMC.jl

-function Hamiltonian(metric::AbstractMetric, ℓ::LogDensityModel; kwargs...)
-    return Hamiltonian(metric, ℓ.logdensity; kwargs...)
+function Hamiltonian(
+    metric::AbstractMetric, ℓ::LogDensityModel; adtype=AutoForwardDiff(), kwargs...


Suggested change

metric::AbstractMetric, ℓ::LogDensityModel; adtype=AutoForwardDiff(), kwargs...

metric::AbstractMetric, ℓ::LogDensityModel; adtype::ADType=AutoForwardDiff(), kwargs...

Suggested change

metric::AbstractMetric, ℓ::LogDensityModel; adtype=AutoForwardDiff(), kwargs...

metric::AbstractMetric, ℓ::LogDensityModel; adtype=ADTypes.AutoForwardDiff(), kwargs...

devmotion · 2025-03-22T21:57:56Z

src/AdvancedHMC.jl

@@ -146,7 +151,12 @@ function Hamiltonian(metric::AbstractMetric, ℓ; kwargs...)
    ℓπ = if cap === LogDensityProblems.LogDensityOrder{0}()
        # In this case ℓ does not support evaluation of the gradient of the log density function
        # We use ForwardDiff to compute the gradient
-        LogDensityProblemsAD.ADgradient(Val(:ForwardDiff), ℓ; kwargs...)


Why not just

Suggested change

LogDensityProblemsAD.ADgradient(Val(:ForwardDiff), ℓ; kwargs...)

LogDensityProblemsAD.ADgradient(adtype, ℓ; kwargs...)

?

devmotion · 2025-03-22T22:01:04Z

src/AdvancedHMC.jl

+        _logdensity_and_gradient = function (x)
+            prep = DI.prepare_gradient(_logdensity, adtype, x)


What's the point of the preparation step in this function? Every time you'd evaluate logdensity + gradient for an input x, it would run preparation for x and then immediately discard it after the computation. If you just use LogDensityProblemsAD.ADgradient (see above), then users could provide a typical input x as keyword argument to optimize the AD computations of the hamiltonian (or otherwise they'd get the slower fallback).

If preparation is not reused, then it is not beneficial. But as per my comment above, if it can be reused because we have access to a typical input, then it works basically the same as a LogDensityProblemsAD.ADgradient (except that Turing no longer depends on LogDensityProblemsAD` in order to mutualize backend code maintenance inside DI)

Changing to DifferentiationInterface is just a step forward for the mutualized AD maintenance in Turing, and seems if we want to use DI to compute the gradient, we will have to use the preparation mechanism. As for the performance aspects, I will test the difference between using DifferentiationInterface and LogDensityProblemsAD, if no difference between those two options, we can safely change to DI then

You don't actually have to use preparation with DI.
DI.gradient(f, backend, x) will actually work too, it usually defaults to preparing first and executing second, but for some backends there are faster shortcuts which we try to take.

devmotion · 2025-03-22T22:01:31Z

src/AdvancedHMC.jl

-function Hamiltonian(
-    metric::AbstractMetric, ℓπ::LogDensityModel, kind::Union{Symbol,Val,Module}; kwargs...
-)


No need to remove these?

yebai · 2025-03-24T19:31:05Z

@ErikQQY, let's only introduce ADTypes support in this PR and switch to DifferentiationInterface via a separate PR.

ErikQQY · 2025-03-29T03:18:06Z

As per @yebai's suggestion, let's only focus on ADTypes support in this PR and integrate DifferentiationInterface in another PR

ext/AdvancedHMCADTypesExt.jl

ErikQQY and others added 2 commits March 17, 2025 01:51

Allow ADTypes AD backend selection in Hamiltonian

e009707

Merge branch 'main' into qqy/adtypes

2a0a36f

yebai requested review from yebai, sunxd3 and devmotion March 17, 2025 12:08

devmotion reviewed Mar 17, 2025

View reviewed changes

yebai assigned ErikQQY Mar 17, 2025

ErikQQY and others added 4 commits March 18, 2025 02:20

Apply reviews

9ae9914

Merge branch 'main' into qqy/adtypes

67c8756

Merge branch 'main' into qqy/adtypes

e5c5a76

Fix CI failings

6d334c8

ErikQQY requested a review from devmotion March 18, 2025 07:19

yebai reviewed Mar 18, 2025

View reviewed changes

ext/AdvancedHMCADTypesExt.jl Show resolved Hide resolved

yebai reviewed Mar 18, 2025

View reviewed changes

test/demo.jl Outdated Show resolved Hide resolved

Integrate DifferentiationInterface

c1c7856

ErikQQY requested a review from yebai March 22, 2025 07:30

gdalle suggested changes Mar 22, 2025

View reviewed changes

Fix typos and remove LDPAD from deps

8cbbc3a

devmotion reviewed Mar 22, 2025

View reviewed changes

ErikQQY added 5 commits March 29, 2025 10:30

Lets only focus on adtypes

abb67eb

Fix some typos

1060e7d

Remove default AD choose

3f33d3c

Misc

6e1af0d

Also test when provided as LogDensityModel

0fdf9ac

devmotion reviewed Mar 29, 2025

View reviewed changes

ext/AdvancedHMCADTypesExt.jl Outdated Show resolved Hide resolved

ErikQQY added 3 commits March 30, 2025 21:00

Dont forget other dependencies for ADTypes ext

0958515

Proper using

7b63145

Proper using

a89cf2f

devmotion approved these changes Mar 30, 2025

View reviewed changes

yebai merged commit 7ea27f3 into TuringLang:main Mar 30, 2025
17 checks passed

ErikQQY mentioned this pull request Apr 1, 2025

Integrate DifferentiationInterface #416

Open

Provide feedback

Saved searches

Use saved searches to filter your results more quickly

Allow ADTypes AD backend selection in Hamiltonian #405

Allow ADTypes AD backend selection in Hamiltonian #405

ErikQQY commented Mar 16, 2025

gdalle Mar 22, 2025

ErikQQY Mar 22, 2025

gdalle Mar 22, 2025

ErikQQY Mar 22, 2025

gdalle Mar 22, 2025

gdalle Mar 22, 2025

devmotion Mar 22, 2025

ErikQQY Mar 23, 2025

devmotion Mar 22, 2025

gdalle Mar 23, 2025

devmotion Mar 22, 2025

gdalle Mar 23, 2025

devmotion Mar 22, 2025

devmotion Mar 22, 2025

gdalle Mar 23, 2025 •

edited

Loading

ErikQQY Mar 23, 2025

gdalle Mar 24, 2025

devmotion Mar 22, 2025

yebai commented Mar 24, 2025

ErikQQY commented Mar 29, 2025

	function Hamiltonian(metric::AbstractMetric, ℓ; adtype=AutoForwardDiff(), kwargs...)
	function Hamiltonian(metric::AbstractMetric, ℓ; adtype::ADType=AutoForwardDiff(), kwargs...)

	metric::AbstractMetric, ℓ::LogDensityModel; adtype=AutoForwardDiff(), kwargs...
	metric::AbstractMetric, ℓ::LogDensityModel; adtype::ADType=AutoForwardDiff(), kwargs...

	LogDensityProblemsAD.ADgradient(Val(:ForwardDiff), ℓ; kwargs...)
	LogDensityProblemsAD.ADgradient(adtype, ℓ; kwargs...)

		_logdensity_and_gradient = function (x)
		prep = DI.prepare_gradient(_logdensity, adtype, x)

Allow ADTypes AD backend selection in Hamiltonian #405

Allow ADTypes AD backend selection in Hamiltonian #405

Conversation

ErikQQY commented Mar 16, 2025

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

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

gdalle Mar 23, 2025 • edited Loading

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

yebai commented Mar 24, 2025

ErikQQY commented Mar 29, 2025

gdalle Mar 23, 2025 •

edited

Loading