Redo the HVP selection

src/DEER/DEER.jl

@@ -189,23 +189,34 @@ statically — `_make_hvp_fn(_hvp_strategy(backend), ...)` resolves to one
 concrete method (and one concrete return type) at compile time, without
 relying on constant propagation through `===`.
 
-Mooncake/Zygote: We route them through
-`ReverseOnGrad` as the default. both CPU and GPU run this path fine (see the
-Mooncake GPU test). `AutoEnzyme` stays on `ForwardOnGrad`: its reverse mode
-hits the gc-transition abort on GPU (see `ext/EnzymeExt.jl`), whereas
-`Enzyme.Forward` is robust on CuArrays once the matmul is wrapped.
-
-Note: forward mode is also supported for these backends via their forward counterparts.
+The choice derives from DI's `pushforward_performance` trait
+mirroring how DI's `hvp_mode` picks its composition:
+backends with a fast pushforward (forward-capable or mode-agnostic, e.g.
+AutoForwardDiff, AutoMooncakeForward, plain AutoEnzyme) take `ForwardOnGrad`;
+reverse-only backends (AutoMooncake, AutoZygote, AutoReverseDiff, AutoTracker,
+AutoEnzyme pinned to Reverse) take `ReverseOnGrad`. Both paths are exercised
+on CPU and GPU (see the Mooncake/Zygote GPU tests). Plain
+`AutoEnzyme()` reports `ForwardOrReverseMode` and thus lands on
+`ForwardOnGrad`. This is fine, since Enzyme's reverse mode hits the 
+gc-transition abort on GPU (see `ext/EnzymeExt.jl`) whereas `Enzyme.Forward` is 
+robust on CuArrays once the matmul is wrapped.
+
+For a `DI.SecondOrder` backend the outer backend differentiates `gradlogp`,
+so the strategy comes from the outer backend's trait.
 =#
 abstract type HVPStrategy end
 struct ForwardOnGrad <: HVPStrategy end
 struct ReverseOnGrad <: HVPStrategy end
 
-_hvp_strategy(::AbstractADType) = ForwardOnGrad()
-_hvp_strategy(::ADTypes.AutoMooncake) = ReverseOnGrad()
-_hvp_strategy(::ADTypes.AutoZygote) = ReverseOnGrad()
-_hvp_strategy(::ADTypes.AutoReverseDiff) = ReverseOnGrad()
-_hvp_strategy(::ADTypes.AutoTracker) = ReverseOnGrad()
+_strategy_from(::DI.PushforwardFast) = ForwardOnGrad()
+_strategy_from(::DI.PushforwardSlow) = ReverseOnGrad()
+
+function _hvp_strategy(backend::AbstractADType)
+    return _strategy_from(DI.pushforward_performance(backend))
+end
+function _hvp_strategy(backend::DI.SecondOrder)
+    return _strategy_from(DI.pushforward_performance(DI.outer(backend)))
+end
 
 #=
 Hooks for backend-specific normalization of the user's `backend`.

src/interface.jl

@@ -372,14 +372,16 @@ function _build_mala_deer_rec(
 
     #=
     Use a model-provided HVP when available. Otherwise compute Hv via AD,
-    picking the path from the user's `backend` (see `DEER._hvp_strategy`):
-
-      ForwardOnGrad() — `pushforward(gradlogp, x, v)`. Used for forward-
-                        capable backends (AutoEnzyme, AutoForwardDiff, AutoMooncakeForward).
-      ReverseOnGrad() — `gradient(x -> pmcmc_dot(gradlogp(x), v))`. Default
-                        routing for AutoMooncake / AutoZygote / AutoReverseDiff
-                        / AutoTracker (which also have forward counterparts,
-                        see `DEER._hvp_strategy` for the nuance).
+    picking the path from the user's `backend` via DI's
+    `pushforward_performance` trait (see `DEER._hvp_strategy`):
+
+      ForwardOnGrad() — `pushforward(gradlogp, x, v)`. Used for backends with
+                        a fast pushforward (AutoEnzyme, AutoForwardDiff,
+                        AutoMooncakeForward).
+      ReverseOnGrad() — `gradient(x -> pmcmc_dot(gradlogp(x), v))`. Used for
+                        reverse-only backends (AutoMooncake / AutoZygote /
+                        AutoReverseDiff), which also have forward counterparts — see 
+                        `DEER._hvp_strategy` for the nuance.
 
     Either path can be bypassed by providing an analytical `hvp` /
     `hvp_batch` on the `DensityModel`.

test/test-HVP-Strategy.jl

@@ -0,0 +1,46 @@
+using Test
+using ADTypes
+using Enzyme: Enzyme
+using ParallelMCMC
+
+const DEER_STRAT = ParallelMCMC.DEER
+const DI_STRAT = ParallelMCMC.DEER.DI
+
+#=
+The AD-HVP fallback strategy is derived from DI's `pushforward_performance`
+trait (see #38): backends with a fast pushforward take ForwardOnGrad, and
+reverse-only backends take ReverseOnGrad — no per-backend enumeration.
+=#
+@testset "HVP strategy from DI mode traits" begin
+    @testset "forward-capable backends → ForwardOnGrad" begin
+        @test DEER_STRAT._hvp_strategy(AutoForwardDiff()) isa DEER_STRAT.ForwardOnGrad
+        @test DEER_STRAT._hvp_strategy(AutoEnzyme()) isa DEER_STRAT.ForwardOnGrad
+        @test DEER_STRAT._hvp_strategy(AutoEnzyme(; mode=Enzyme.Forward)) isa
+            DEER_STRAT.ForwardOnGrad
+        @test DEER_STRAT._hvp_strategy(AutoMooncakeForward()) isa DEER_STRAT.ForwardOnGrad
+    end
+
+    @testset "reverse-only backends → ReverseOnGrad" begin
+        @test DEER_STRAT._hvp_strategy(AutoMooncake()) isa DEER_STRAT.ReverseOnGrad
+        @test DEER_STRAT._hvp_strategy(AutoZygote()) isa DEER_STRAT.ReverseOnGrad
+        @test DEER_STRAT._hvp_strategy(AutoReverseDiff()) isa DEER_STRAT.ReverseOnGrad
+        @test DEER_STRAT._hvp_strategy(AutoTracker()) isa DEER_STRAT.ReverseOnGrad
+        @test DEER_STRAT._hvp_strategy(AutoEnzyme(; mode=Enzyme.Reverse)) isa
+            DEER_STRAT.ReverseOnGrad
+    end
+
+    @testset "SecondOrder picks the strategy from the outer backend" begin
+        so_fwd_outer = DI_STRAT.SecondOrder(AutoForwardDiff(), AutoZygote())
+        @test DEER_STRAT._hvp_strategy(so_fwd_outer) isa DEER_STRAT.ForwardOnGrad
+
+        so_rev_outer = DI_STRAT.SecondOrder(AutoZygote(), AutoForwardDiff())
+        @test DEER_STRAT._hvp_strategy(so_rev_outer) isa DEER_STRAT.ReverseOnGrad
+    end
+
+    @testset "strategy resolution is type-stable" begin
+        @test @inferred(DEER_STRAT._hvp_strategy(AutoForwardDiff())) isa
+            DEER_STRAT.ForwardOnGrad
+        @test @inferred(DEER_STRAT._hvp_strategy(AutoMooncake())) isa
+            DEER_STRAT.ReverseOnGrad
+    end
+end