PR Review for #62

CHANGELOG.md

@@ -50,12 +50,33 @@ and this project adheres to [Semantic Versioning](https://semver.org/spec/v2.0.0
   its strategy was routed on `DI.outer(backend)`, so a
   `DifferentiationInterface.SecondOrder` ran the wrong half of the pair. A
   `SecondOrder` now goes to the true second-order path instead of either
-  strategy, and the half-selecting helpers it still uses agree: normalization
-  applies to the outer, the pass whose mode the backend extensions care about.
-  Unwrapping to that half happens before the normalization hooks are dispatched
-  on, so a `SecondOrder(AutoEnzyme(), ...)` still reaches `EnzymeExt` and gets
-  its mode and function annotation pinned rather than running as a bare
-  `AutoEnzyme()` (which aborts on GPU).
+  strategy, and the half-selecting helper it still uses agrees: normalization
+  applies to the outer pass. Unwrapping to that half happens before the
+  normalization hook is dispatched on, so a `SecondOrder(AutoEnzyme(), ...)` still
+  reaches `EnzymeExt` and gets its function annotation filled in rather than
+  running as a bare `AutoEnzyme()`.
+- Backend normalization no longer picks a differentiation mode on the user's
+  behalf (#62). `EnzymeExt` pinned `mode=Enzyme.Forward` (with
+  `set_runtime_activity`) onto an `AutoEnzyme()` left mode-agnostic, on the
+  grounds that reverse mode hit a gc-transition abort on GPU and that composed
+  `pmcmc_matmul` calls needed runtime activity. The `pmcmc_*` Enzyme rules keep
+  Enzyme off both paths on their own now, so the pin bought nothing — and it cost
+  correctness, because it silently rewrote the direction of a `SecondOrder`'s
+  outer half. `SecondOrder(AutoEnzyme(), AutoForwardDiff())` is
+  reverse-over-forward to `hvp_mode`, its inner half being forward-only, and came
+  out forward-over-forward. Normalization now fills in only
+  `function_annotation=Enzyme.Const`, which is about this package's own read-only
+  HVP wrappers rather than about Enzyme's mode, and leaves `mode` exactly as given
+  — unset included, for DI to resolve from the operator it runs. `hvp_mode` is
+  therefore identical before and after normalization for every backend pair.
+
+  A mode set explicitly was never overridden, so only mode-agnostic backends were
+  affected, and the HVP was a correct HVP either way; what changes is that the
+  composition asked for is the one that runs. The two normalization hooks
+  (`_hvp_forward_backend`, `_hvp_closure_backend`) collapse into a single
+  `_normalized_backend`, since without a mode to choose they no longer differ.
+  Users relying on a plain `AutoEnzyme()` being run forward should now pass
+  `AutoEnzyme(; mode=Enzyme.Forward)` explicitly.
 
 ### Changed
 

ext/EnzymeExt.jl

@@ -20,46 +20,26 @@ using Enzyme.EnzymeCore.EnzymeRules:
 # TODO: Implement matmul overloads upstream in Enzyme. See: https://github.com/EnzymeAD/Enzyme.jl/issues/3122
 
 #=
-Tell DEER's forward-on-grad HVP path how to normalize a plain `AutoEnzyme()`:
-pin `mode=Enzyme.Forward` and `function_annotation=Enzyme.Const`. Pinning
-Forward is load-bearing on GPU — without it DI defaults to reverse mode, which
-hits the gc-transition abort documented on the Enzyme rules below.
+Normalization of the user's `AutoEnzyme` for DEER's AD-HVP paths: fill in
+`function_annotation=Enzyme.Const` when they left it open, so Enzyme doesn't
+throw `EnzymeMutabilityException` on the read-only `_HvpReverseClosure` /
+`_BatchHvpReverseClosure` wrappers, which capture `gradlogp`. Those wrapper types
+belong to this package, so declaring them constant is this package's business.
 
-`mode` and `function_annotation` are normalized independently — a user who
-sets one keeps that choice, but still gets the default for the other.
+`mode` is passed through exactly as given, unset included. Choosing a direction
+on the user's behalf is not our call: a mode they set is a decision, and an unset
+one is DI's to resolve from the operator it runs.
 
-`set_runtime_activity` is load-bearing for composed `pmcmc_matmul` calls
-(e.g. `pmcmc_matmul(transpose(X), pmcmc_matmul(X, β))`). Static activity
-analysis can't prove the outer call's `transpose(X)` shadow is safe to
-reuse, and Enzyme aborts with `EnzymeRuntimeActivityError`. With runtime
-activity, the shadow is tracked dynamically.
+An earlier version pinned `mode=Enzyme.Forward` here (with
+`set_runtime_activity`) against the gc-transition abort on GPU and
+`EnzymeRuntimeActivityError` on composed `pmcmc_matmul` calls. The rules below
+keep Enzyme off both paths on their own, so the pin bought nothing and cost
+correctness: it silently rewrote the direction of a `SecondOrder`'s outer half
+(see `DEER._normalized_backend`).
 =#
-function DEER._hvp_forward_backend(backend::ADTypes.AutoEnzyme{M,A}) where {M,A}
-    mode = if backend.mode === nothing
-        Enzyme.set_runtime_activity(Enzyme.Forward)
-    else
-        backend.mode
-    end
-    annotation = A === Nothing ? Enzyme.Const : A
-    return ADTypes.AutoEnzyme(; mode=mode, function_annotation=annotation)
-end
-
-#=
-Tell DEER's reverse-on-grad HVP path how to normalize a plain `AutoEnzyme()`:
-fill in `function_annotation=Enzyme.Const` so Enzyme doesn't throw
-`EnzymeMutabilityException` on the read-only `_HvpReverseClosure` /
-`_BatchHvpReverseClosure` wrappers, and default `mode` to reverse with
-runtime activity. As in `_hvp_forward_backend`, the two fields are
-normalized independently.
-=#
-function DEER._hvp_closure_backend(backend::ADTypes.AutoEnzyme{M,A}) where {M,A}
-    mode = if backend.mode === nothing
-        Enzyme.set_runtime_activity(Enzyme.Reverse)
-    else
-        backend.mode
-    end
-    annotation = A === Nothing ? Enzyme.Const : A
-    return ADTypes.AutoEnzyme(; mode=mode, function_annotation=annotation)
+function DEER._normalized_backend(backend::ADTypes.AutoEnzyme{M,A}) where {M,A}
+    A === Nothing || return backend
+    return ADTypes.AutoEnzyme(; mode=backend.mode, function_annotation=Enzyme.Const)
 end
 
 #=

src/DEER/DEER.jl

@@ -93,7 +93,7 @@ function _prepare_hvp(f, backend::AbstractADType, x_template::AbstractVector)
     v_template = similar(x_template)
     fill!(v_template, zero(eltype(x_template)))
     return DI.prepare_pushforward(
-        f, _hvp_forward_backend(backend), x_template, (v_template,)
+        f, _normalized_backend(backend), x_template, (v_template,)
     )
 end
 
@@ -116,14 +116,14 @@ function _hvp_prepared(
 )
     x_exec = _materialize_ad_vector(x)
     v_exec = _tangent_like(x_exec, v)
-    res = DI.pushforward(f, prep, _hvp_forward_backend(backend), x_exec, (v_exec,))
+    res = DI.pushforward(f, prep, _normalized_backend(backend), x_exec, (v_exec,))
     return res isa Tuple ? first(res) : res
 end
 
 function _hvp_nopre(f, backend::AbstractADType, x::AbstractVector, v::AbstractVector)
     x_exec = _materialize_ad_vector(x)
     v_exec = _tangent_like(x_exec, v)
-    res = DI.pushforward(f, _hvp_forward_backend(backend), x_exec, (v_exec,))
+    res = DI.pushforward(f, _normalized_backend(backend), x_exec, (v_exec,))
     return res isa Tuple ? first(res) : res
 end
 
@@ -133,7 +133,7 @@ function _prepare_batch_hvp_from_grad(
     V_template = similar(X_template)
     fill!(V_template, zero(eltype(X_template)))
     return DI.prepare_pushforward(
-        grad_batch, _hvp_forward_backend(backend), X_template, (V_template,)
+        grad_batch, _normalized_backend(backend), X_template, (V_template,)
     )
 end
 
@@ -142,7 +142,7 @@ function _batch_hvp_from_grad_prepared(
 )
     X_exec = _materialize_ad_matrix(X)
     V_exec = _tangent_like(X_exec, V)
-    res = DI.pushforward(grad_batch, prep, _hvp_forward_backend(backend), X_exec, (V_exec,))
+    res = DI.pushforward(grad_batch, prep, _normalized_backend(backend), X_exec, (V_exec,))
     return res isa Tuple ? first(res) : res
 end
 
@@ -196,9 +196,8 @@ relying on constant propagation through `===`.
 The routing follows DI's `hvp_mode`: a forward outer pass
 (`DI.ForwardOverAnything`) takes `ForwardOnGrad`, anything else
 `ReverseOnGrad`. Only the outer direction matters since we differentiate
-the already-built `gradlogp`. Plain `AutoEnzyme()` lands on `ForwardOnGrad`,
-which we need: Enzyme reverse hits the gc-transition abort on GPU (see
-`ext/EnzymeExt.jl`).
+the already-built `gradlogp`. The mode is whatever the user's backend carries;
+nothing here substitutes one (see `_normalized_backend`).
 =#
 abstract type HVPStrategy end
 struct ForwardOnGrad <: HVPStrategy end
@@ -207,46 +206,39 @@ struct ReverseOnGrad <: HVPStrategy end
 _strategy_from(::DI.ForwardOverAnything) = ForwardOnGrad()
 _strategy_from(::DI.HVPMode) = ReverseOnGrad()
 
-function _hvp_strategy(backend::Union{AbstractADType,DI.SecondOrder})
-    return _strategy_from(DI.hvp_mode(backend))
-end
+_hvp_strategy(backend::AbstractADType) = _strategy_from(DI.hvp_mode(backend))
 
 #=
-Hooks for backend-specific normalization of the user's `backend`.
-
-`_hvp_forward_backend` is for the forward-on-grad pushforward path
-(differentiates the user's `gradlogp` directly). EnzymeExt specializes it
-to pin `mode=Enzyme.Forward` and `function_annotation=Enzyme.Const` when
-the user passed plain `AutoEnzyme()`, without pinning Forward, DI lowers
-through reverse mode and hits the gc-transition abort on GPU (see
-`ext/EnzymeExt.jl`).
-
-`_hvp_closure_backend` is for the reverse-on-grad gradient path on the
-read-only `_HvpReverseClosure` / `_BatchHvpReverseClosure` wrappers.
-EnzymeExt specializes it to set `function_annotation=Enzyme.Const` so
-Enzyme doesn't throw `EnzymeMutabilityException` on a closure that captures
-`gradlogp`.
-
-A `SecondOrder` is unwrapped to its outer half first, since the pass we are
-about to run is the outer one i.e., the inner derivative is whatever
-`gradlogp` already is. That keeps the half `_hvp_strategy` routed on, so the
-strategy and the backend that carries it out can't end up disagreeing. The
-unwrapping recurses rather than calling `DI.outer` in the generic method, so
-that a wrapped backend still reaches its own normalization: dispatch happens
-on what comes out of `DI.outer`, not on the `SecondOrder` around it.
+Hook for backend-specific normalization of the user's `backend`, applied on every
+AD-HVP path before the backend reaches DI.
+
+It supplies what the wrappers DEER differentiates need, and nothing else. Those
+wrapper types are ours, so annotating them is ours to do: EnzymeExt specializes
+this to fill `function_annotation=Enzyme.Const`, without which Enzyme throws
+`EnzymeMutabilityException` on the read-only `_HvpReverseClosure` /
+`_BatchHvpReverseClosure`, which capture `gradlogp`.
+
+It deliberately does not choose a differentiation mode. Which direction a pass
+runs is the user's call when they state one and DI's to resolve from the operator
+when they don't; this package is not an AD package and has no business overriding
+either. Picking one here also used to corrupt a `SecondOrder`, whose halves carry
+directions of their own (see `_normalized_second_order`).
+
+Callers hand this a single pass, never a `SecondOrder`: the strategy paths below
+run one AD pass over a hand-written `gradlogp`, and `_resolve_hvp` sends every
+`SecondOrder` to `_make_hvp_fn_second_order` before they are reached.
+`_normalized_second_order` is the one caller that starts from a pair, and it
+selects the outer half itself.
 =#
-_hvp_forward_backend(backend::DI.SecondOrder) = _hvp_forward_backend(DI.outer(backend))
-_hvp_closure_backend(backend::DI.SecondOrder) = _hvp_closure_backend(DI.outer(backend))
-_hvp_forward_backend(backend::AbstractADType) = backend
-_hvp_closure_backend(backend::AbstractADType) = backend
+_normalized_backend(backend::AbstractADType) = backend
 
 function _prepare_hvp_via_grad_reverse(
     gradlogp, backend::AbstractADType, x_template::AbstractVector
 )
     v_template = similar(x_template)
     fill!(v_template, zero(eltype(x_template)))
     f = _HvpReverseClosure(gradlogp)
-    eff_backend = _hvp_closure_backend(backend)
+    eff_backend = _normalized_backend(backend)
     prep = DI.prepare_gradient(f, eff_backend, x_template, DI.Constant(v_template))
     return (f, prep, eff_backend)
 end
@@ -262,7 +254,7 @@ function _prepare_batch_hvp_via_grad_reverse(
     V_template = similar(X_template)
     fill!(V_template, zero(eltype(X_template)))
     f = _BatchHvpReverseClosure(grad_batch)
-    eff_backend = _hvp_closure_backend(backend)
+    eff_backend = _normalized_backend(backend)
     prep = DI.prepare_gradient(f, eff_backend, X_template, DI.Constant(V_template))
     return (f, prep, eff_backend)
 end
@@ -315,16 +307,20 @@ takes both passes over the log-density, so these never touch the gradient slot.
 Preferred over pushing tangents through a prepared DI gradient, which drops out
 of its preparation once the outer pass hands it an unexpected tangent type.
 
-Only the outer half is normalized: `_hvp_forward_backend` selects it out of the
-pair. The inner is a plain first-order gradient and needs no pinning.
+Both halves are passed to `DI.hvp` as the user composed them, so the direction
+each one runs in is theirs and DI's, not ours. Normalization touches only the
+outer half, and only to fill in annotations for the wrappers being
+differentiated; because it never substitutes a mode, `DI.hvp_mode` of the pair is
+the same before and after. The inner half is a plain first-order gradient over
+the user's own `logdensity` and is passed straight through.
 
 The batched form differentiates `sum(logdensity_batch(X))`, whose Hessian is
 block-diagonal by column independence, so its HVP along `V` is the columnwise
 HVP. Same argument the batched gradient rests on.
 ---------------------------------------------------------------------------
 =#
 function _normalized_second_order(backend::DI.SecondOrder)
-    return DI.SecondOrder(_hvp_forward_backend(backend), DI.inner(backend))
+    return DI.SecondOrder(_normalized_backend(DI.outer(backend)), DI.inner(backend))
 end
 
 function _make_hvp_fn_second_order(

src/interface.jl

@@ -27,8 +27,8 @@ DI's second-order operator. Passing a `SecondOrder` yourself always means the
 latter, and bypasses the gradient slot even when you wrote it by hand.
 
 - `logdensity(x::AbstractVector) -> Real`
-- `grad_logdensity` — callable `x -> AbstractVector`, or a backend applied to
-  `logdensity`.
+- `grad_logdensity` — callable `x -> AbstractVector`, or a backend to
+  differentiate `logdensity` with.
 - `hvp` — optional callable `(x, v) -> AbstractVector`, or a backend. If
   `nothing`, DEER builds the HVP from the sampler's `backend`.
 - `logdensity_batch(X::AbstractMatrix) -> AbstractVector` — optional batched
@@ -37,7 +37,7 @@ latter, and bypasses the gradient slot even when you wrote it by hand.
   A batched gradient derived from this is one gradient of its sum, so coupling
   between columns would go unnoticed and give wrong derivatives.
 - `grad_logdensity_batch` — optional callable `X -> AbstractMatrix`, or a
-  backend applied to `logdensity_batch`. Left out alongside a
+  backend to differentiate `logdensity_batch` with. Left out alongside a
   `logdensity_batch`, it is derived when `grad_logdensity` is a backend.
 - `hvp_batch` — optional callable `(X, V) -> AbstractMatrix`, or a backend,
   resolved against `grad_logdensity_batch` the same way `hvp` is against

test/test-HVP-Strategy.jl

@@ -36,26 +36,56 @@ const DI_STRAT = ParallelMCMC.DEER.DI
         @test DEER_STRAT._hvp_strategy(so_agnostic_outer) isa DEER_STRAT.ReverseOnGrad
     end
 
-    @testset "the backend that runs is the one routed on" begin
-        # both paths differentiate the already-built gradlogp, so both take the outer
-        so_fwd = DI_STRAT.SecondOrder(AutoForwardDiff(), AutoZygote())
-        @test DEER_STRAT._hvp_forward_backend(so_fwd) === AutoForwardDiff()
+    @testset "normalization supplies Const but never a mode" begin
+        #= The wrappers DEER differentiates are its own types, so annotating them
+        `Const` is its business. The mode is not: one the user set is a decision,
+        and an unset one is DI's to resolve from the operator it runs. =#
+        bare = DEER_STRAT._normalized_backend(AutoEnzyme())
+        @test bare isa AutoEnzyme{<:Any,Enzyme.Const}
+        @test bare.mode === nothing
 
-        so_rev = DI_STRAT.SecondOrder(AutoZygote(), AutoForwardDiff())
-        @test DEER_STRAT._hvp_closure_backend(so_rev) === AutoZygote()
+        for mode in (Enzyme.Forward, Enzyme.Reverse)
+            normalized = DEER_STRAT._normalized_backend(AutoEnzyme(; mode=mode))
+            @test normalized.mode === mode
+            @test normalized isa AutoEnzyme{<:Any,Enzyme.Const}
+        end
+
+        # An annotation the user chose is left alone.
+        annotated = AutoEnzyme(; function_annotation=Enzyme.Duplicated)
+        @test DEER_STRAT._normalized_backend(annotated) === annotated
+
+        # Backends with no specialization pass straight through.
+        @test DEER_STRAT._normalized_backend(AutoForwardDiff()) === AutoForwardDiff()
+        @test DEER_STRAT._normalized_backend(AutoZygote()) === AutoZygote()
     end
 
-    @testset "unwrapping a SecondOrder still reaches backend normalization" begin
-        #= The outer half has to be taken before the backend-specific hook is
-        dispatched on, or a wrapped `AutoEnzyme()` comes out bare: unnormalized,
-        it lowers through reverse mode and aborts on GPU (see ext/EnzymeExt.jl). =#
-        so = DI_STRAT.SecondOrder(AutoEnzyme(), AutoForwardDiff())
-        @test DEER_STRAT._hvp_forward_backend(so) ===
-            DEER_STRAT._hvp_forward_backend(AutoEnzyme())
-        @test DEER_STRAT._hvp_closure_backend(so) ===
-            DEER_STRAT._hvp_closure_backend(AutoEnzyme())
-        @test DEER_STRAT._hvp_forward_backend(so).mode isa Enzyme.ForwardMode
-        @test DEER_STRAT._hvp_closure_backend(so) isa AutoEnzyme{<:Any,Enzyme.Const}
+    @testset "normalizing a SecondOrder keeps the composition DI resolved" begin
+        #= Regression for #62. Normalization used to route the outer half through
+        a forward-only hook, which pinned `Enzyme.Forward` onto it. For a pair
+        `hvp_mode` resolves to reverse -- `SecondOrder(AutoEnzyme(),
+        AutoForwardDiff())` is reverse-over-forward, its inner half being
+        forward-only -- that silently made it forward-over-forward. =#
+        for so in (
+            DI_STRAT.SecondOrder(AutoEnzyme(), AutoForwardDiff()),
+            DI_STRAT.SecondOrder(AutoEnzyme(), AutoZygote()),
+            DI_STRAT.SecondOrder(AutoEnzyme(; mode=Enzyme.Reverse), AutoForwardDiff()),
+            DI_STRAT.SecondOrder(AutoEnzyme(; mode=Enzyme.Forward), AutoZygote()),
+            DI_STRAT.SecondOrder(AutoForwardDiff(), AutoZygote()),
+            DI_STRAT.SecondOrder(AutoZygote(), AutoForwardDiff()),
+        )
+            normalized = DEER_STRAT._normalized_second_order(so)
+            @test DI_STRAT.hvp_mode(normalized) == DI_STRAT.hvp_mode(so)
+            # The inner half is the user's own first-order gradient, untouched.
+            @test DI_STRAT.inner(normalized) === DI_STRAT.inner(so)
+            # And no mode is invented for the outer half either. Only `AutoEnzyme`
+            # carries a `mode`, so this is the case that could regress.
+            if DI_STRAT.outer(so) isa AutoEnzyme
+                @test DI_STRAT.outer(normalized).mode === DI_STRAT.outer(so).mode
+                #= The outer half still reaches EnzymeExt's specialization rather
+                than passing through as a bare backend. =#
+                @test DI_STRAT.outer(normalized) isa AutoEnzyme{<:Any,Enzyme.Const}
+            end
+        end
     end
 
     @testset "strategy resolution is type-stable" begin