Skip to content

Commit 4390c43

Browse files
yebaiclaude
yebai
and
claude
authored
Add order=2 Hessian preparation via value_gradient_and_hessian!! (#163)
* Add order=2 Hessian preparation via value_gradient_and_hessian!! Extends `prepare(adtype, problem, x; order=2)` to build Hessian machinery for scalar-valued problems on the DI and Mooncake extensions, returning `(value, gradient, hessian)` from a new `value_gradient_and_hessian!!` generic. Unifies the per-extension caches (`DICache`, `MooncakeCache`) so one struct carries every derivative order, with explicit cross-arity error messages replacing prior `MethodError`s. DI uses the in-place `DI.value_gradient_and_hessian!` with caller-owned buffers; Mooncake uses its native `prepare_hessian_cache` API. Co-Authored-By: Claude Opus 4.7 (1M context) <noreply@anthropic.com> * Split Hessian edge case into separate :hessian_edge testcase group Move the order=1-prep error case for value_gradient_and_hessian!! out of :edge and into a new :hessian_edge group so Hessian-specific edge checks are only exercised by preparations that support order=2. Co-Authored-By: Claude Opus 4.7 (1M context) <noreply@anthropic.com> * Update Project.toml * Encode derivative order on Prepared; split DICache into per-arity types `Prepared` gains an `Order` type parameter (`1` for gradient/jacobian, `2` for Hessian) with an `order(::Prepared)` accessor, so the prep order can be retrieved reliably without inspecting the backend-specific cache type. `value_and_gradient!!` on an order=2 prep now returns `(value, gradient)` via a dedicated gradient prep built alongside the Hessian prep — no O(n²) Hessian work for a gradient-only call. For DI's `SecondOrder` backend the gradient prep uses `DI.inner(adtype)` per DI's convention; the same unwrap runs on the hot path so prep and call use matching adtypes. `order` is now validated up-front via `Evaluators._validate_ad_order` (was duplicated across both extensions and fired only after the structural prep had already called `problem` once). DI: `DICache` is replaced by three concrete types — `DIGradientCache`, `DIJacobianCache`, `DIHessianCache` — eliminating the 6-nullable-field struct and runtime `=== nothing` checks. `_di_call_shape` is the shared target-and-constants helper used by both `_prepare_di` (order=1) and the order=2 path; the two preps share one target instance so compiled-tape ReverseDiff sees a consistent `Fix2` closure. Mooncake: `MooncakeCache` gains a `gradient_cache` field populated only at order=2; `_mooncake_gradient_cache` is now used by the NamedTuple path, the order=1 scalar branch, and the order=2 gradient prep. Co-Authored-By: Claude Opus 4.7 (1M context) <noreply@anthropic.com> * Apply JuliaFormatter v1 to changes from previous commit Co-Authored-By: Claude Opus 4.7 (1M context) <noreply@anthropic.com> --------- Co-authored-by: Claude Opus 4.7 (1M context) <noreply@anthropic.com>
1 parent 1599516 commit 4390c43

9 files changed

Lines changed: 614 additions & 138 deletions

File tree

‎Project.toml‎

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -3,7 +3,7 @@ uuid = "7a57a42e-76ec-4ea3-a279-07e840d6d9cf"
33
keywords = ["probabilistic programming"]
44
license = "MIT"
55
desc = "Common interfaces for probabilistic programming"
6-
version = "0.15"
6+
version = "0.15.1"
77

88
[deps]
99
ADTypes = "47edcb42-4c32-4615-8424-f2b9edc5f35b"

‎docs/src/evaluators.md‎

Lines changed: 52 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -138,6 +138,57 @@ library invokes the inner callable many times with same-length dual arrays
138138
derived from a single user-supplied `x`; re-validating on each invocation
139139
would be redundant work in the hot path.
140140

141+
## Hessian (`order=2`)
142+
143+
Pass `order=2` to `prepare` to build a Hessian-capable evaluator. The
144+
returned object answers `value_gradient_and_hessian!!`, which returns
145+
`(value, gradient, hessian)` in a single call. `order=2` requires
146+
`problem` to be scalar-valued; a vector-valued probe throws at preparation
147+
time.
148+
149+
```julia
150+
using AbstractPPL: prepare, value_gradient_and_hessian!!
151+
using ADTypes: AutoForwardDiff
152+
using ForwardDiff, DifferentiationInterface
153+
154+
quadratic(x) = sum(abs2, x)
155+
prepared = prepare(AutoForwardDiff(), quadratic, zeros(3); order=2)
156+
val, grad, hess = value_gradient_and_hessian!!(prepared, [1.0, 2.0, 3.0])
157+
# val == 14.0
158+
# grad == [2.0, 4.0, 6.0]
159+
# hess == [2 0 0; 0 2 0; 0 0 2]
160+
```
161+
162+
Both `context=` and `check_dims=` apply to `order=2` preps with the same
163+
semantics as for `order=1`. The `!!` aliasing contract also extends: the
164+
returned gradient and Hessian may alias internal cache buffers of
165+
`prepared`, so copy before retaining them past the next call. NamedTuple
166+
inputs are not supported at `order=2`.
167+
168+
For DifferentiationInterface, `adtype` can be either a single backend
169+
(letting DI pick its own Hessian strategy) or a
170+
[`DifferentiationInterface.SecondOrder(outer, inner)`](https://juliadiff.org/DifferentiationInterface.jl/stable/api/#DifferentiationInterface.SecondOrder)
171+
composition that selects the outer differentiator and the inner gradient
172+
backend independently — typically forward-over-reverse:
173+
174+
```julia
175+
using DifferentiationInterface: SecondOrder
176+
using ADTypes: AutoForwardDiff, AutoReverseDiff
177+
178+
adtype = SecondOrder(AutoForwardDiff(), AutoReverseDiff())
179+
prepared = prepare(adtype, quadratic, zeros(3); order=2)
180+
```
181+
182+
`SecondOrder <: AbstractADType`, so the same `prepare(adtype, problem, x; order=2)`
183+
entry handles it.
184+
185+
Calling `value_gradient_and_hessian!!` on an `order=1` prep throws an
186+
`ArgumentError` — re-prepare with `order=2` instead. The reverse is allowed:
187+
`value_and_gradient!!` on an `order=2` prep returns `(value, gradient)`
188+
without paying the Hessian cost, since `prepare` builds a dedicated
189+
gradient prep alongside the Hessian one. `value_and_jacobian!!` is rejected
190+
because `order=2` requires a scalar-valued problem.
191+
141192
## Constant context arguments
142193

143194
When the underlying callable naturally takes the form `f(x, context...)` —
@@ -177,4 +228,5 @@ p([1.0, 2.0, 3.0])
177228
AbstractPPL.prepare
178229
AbstractPPL.value_and_gradient!!
179230
AbstractPPL.value_and_jacobian!!
231+
AbstractPPL.value_gradient_and_hessian!!
180232
```

‎ext/AbstractPPLDifferentiationInterfaceExt.jl‎

Lines changed: 193 additions & 60 deletions
Original file line numberDiff line numberDiff line change
@@ -5,27 +5,55 @@ using AbstractPPL.Evaluators: Evaluators, Prepared, VectorEvaluator, _ad_output_
55
using ADTypes: AbstractADType, AutoReverseDiff
66
using DifferentiationInterface: DifferentiationInterface as DI
77

8-
# AD target used by both `DICache` modes. `Vararg{Any,N}` with a free `N`
9-
# forces specialization on the trailing arity (a bare `Vararg{Any}` would
10-
# skip it). DI invokes this as `_call_evaluator(x, f, c1, …, cN)` on the
11-
# constants path, and as `_call_evaluator(x, evaluator)` (via `Fix2`) on
12-
# the closure path — empty `ctx` then makes the splat a no-op.
8+
# AD target used by every DI cache. `Vararg{Any,N}` with a free `N` forces
9+
# specialization on the trailing arity (a bare `Vararg{Any}` would skip it).
10+
# DI invokes this as `_call_evaluator(x, f, c1, …, cN)` on the constants path,
11+
# and as `_call_evaluator(x, evaluator)` (via `Fix2`) on the closure path —
12+
# empty `ctx` then makes the splat a no-op.
1313
@inline _call_evaluator(x, f::F, ctx::Vararg{Any,N}) where {F,N} = f(x, ctx...)
1414

15-
# `Mode` tags the cache shape:
16-
# * `:closure` — compiled-tape ReverseDiff: target is a `Fix2` closure,
17-
# the AD call passes **0** `DI.Constant`s.
18-
# * `N::Int` — constants path: `N == length(evaluator.context)`, the
19-
# AD call passes **N + 1** `DI.Constant`s (`f` plus the
20-
# `N` context values).
21-
# Encoding `Mode` in the type resolves the dispatch in `_di_value_and_*`
22-
# at compile time without a runtime branch.
23-
struct DICache{Mode,F,GP,JP}
15+
# `Mode` tags the call shape:
16+
# * `:closure` — compiled-tape ReverseDiff: target is a `Fix2` closure; the
17+
# AD call passes **0** `DI.Constant`s.
18+
# * `N::Int` — constants path: `N == length(evaluator.context)`; the AD
19+
# call passes **N + 1** `DI.Constant`s (`f` plus the `N`
20+
# context values).
21+
# Encoding `Mode` in each cache type resolves the closure-vs-constants dispatch
22+
# in `_di_value_and_*` at compile time without a runtime branch.
23+
24+
# `Nothing` in the prep slot flags the empty-input cache (DI prep paths fail
25+
# on length-0 input, e.g. ForwardDiff `BoundsError`). Hot paths dispatch on the
26+
# `Nothing` parameter to short-circuit before any DI call. Same convention for
27+
# `DIJacobianCache` and `DIHessianCache` below.
28+
struct DIGradientCache{Mode,F,GP}
2429
target::F
2530
gradient_prep::GP
31+
function DIGradientCache(target::F, gp::GP, ::Val{Mode}) where {Mode,F,GP}
32+
return new{Mode,F,GP}(target, gp)
33+
end
34+
end
35+
36+
struct DIJacobianCache{Mode,F,JP}
37+
target::F
2638
jacobian_prep::JP
27-
function DICache{Mode}(target::F, gp::GP, jp::JP) where {Mode,F,GP,JP}
28-
return new{Mode,F,GP,JP}(target, gp, jp)
39+
function DIJacobianCache(target::F, jp::JP, ::Val{Mode}) where {Mode,F,JP}
40+
return new{Mode,F,JP}(target, jp)
41+
end
42+
end
43+
44+
# Order=2 (scalar-output). `grad_buf` / `hess_buf` are caller-owned output
45+
# buffers handed to `DI.value_gradient_and_hessian!`; the returned arrays alias
46+
# them (`!!` contract).
47+
struct DIHessianCache{Mode,F,GP,HP,G,H}
48+
target::F
49+
gradient_prep::GP
50+
hessian_prep::HP
51+
grad_buf::G
52+
hess_buf::H
53+
function DIHessianCache(
54+
target::F, gp::GP, hp::HP, g::G, h::H, ::Val{Mode}
55+
) where {Mode,F,GP,HP,G,H}
56+
return new{Mode,F,GP,HP,G,H}(target, gp, hp, g, h)
2957
end
3058
end
3159

@@ -34,69 +62,107 @@ end
3462
# target and call DI without constants. Context (if any) is captured inside
3563
# the evaluator closure rather than lowered out — the lowered path would also
3664
# require a closure here, so the wrapper cost is unavoidable for compiled tapes.
37-
function _prepare_di(prep::F, adtype::AutoReverseDiff{true}, x, evaluator) where {F}
38-
target = Base.Fix2(_call_evaluator, evaluator)
39-
return target, prep(target, adtype, x), Val(:closure)
65+
#
66+
# `_di_call_shape` returns `(target, mode, constants)`. For the closure path
67+
# `constants == ()` and the splat at every prep/call site collapses to nothing,
68+
# letting prep and call sites share one shape regardless of mode.
69+
function _di_call_shape(::AutoReverseDiff{true}, evaluator)
70+
return Base.Fix2(_call_evaluator, evaluator), Val(:closure), ()
4071
end
41-
42-
function _prepare_di(prep::F, adtype::AbstractADType, x, evaluator) where {F}
43-
constants = (DI.Constant(evaluator.f), map(DI.Constant, evaluator.context)...)
44-
return (
45-
_call_evaluator,
46-
prep(_call_evaluator, adtype, x, constants...),
47-
Val(length(evaluator.context)),
48-
)
72+
function _di_call_shape(::AbstractADType, evaluator)
73+
return _call_evaluator,
74+
Val(length(evaluator.context)),
75+
(DI.Constant(evaluator.f), map(DI.Constant, evaluator.context)...)
4976
end
5077

51-
@inline _wrap_cache(target, gp, jp, ::Val{Mode}) where {Mode} =
52-
DICache{Mode}(target, gp, jp)
78+
# `SecondOrder` doesn't define gradient prep; per DI's contract the inner
79+
# adtype is the one used for the first derivative.
80+
@inline _gradient_adtype(adtype::AbstractADType) = adtype
81+
@inline _gradient_adtype(adtype::DI.SecondOrder) = DI.inner(adtype)
82+
83+
function _prepare_di(prep::F, adtype, x, evaluator) where {F}
84+
target, mode, constants = _di_call_shape(adtype, evaluator)
85+
return target, prep(target, adtype, x, constants...), mode
86+
end
5387

5488
function AbstractPPL.prepare(
5589
adtype::AbstractADType,
5690
problem,
5791
x::AbstractVector{<:Real};
5892
check_dims::Bool=true,
5993
context::Tuple=(),
94+
order::Int=1,
6095
)
96+
Evaluators._validate_ad_order(order)
6197
evaluator = AbstractPPL.prepare(problem, x; check_dims, context)::VectorEvaluator
6298
arity = _ad_output_arity(evaluator(x))
99+
mode_empty = Val(length(context))
100+
if order == 2
101+
arity === :scalar || Evaluators._throw_hessian_needs_scalar()
102+
if length(x) == 0
103+
cache = DIHessianCache(
104+
_call_evaluator, nothing, nothing, nothing, nothing, mode_empty
105+
)
106+
return Prepared(adtype, evaluator, cache, Val(2))
107+
end
108+
# Build both gradient and Hessian preps against the same target so
109+
# `value_and_gradient!!` on the order=2 prep skips the O(n²) Hessian
110+
# cost. Sharing the target matters for compiled-tape ReverseDiff —
111+
# two `Fix2` instances may not be interchangeable in DI.
112+
target, mode, constants = _di_call_shape(adtype, evaluator)
113+
gradient_prep = DI.prepare_gradient(
114+
target, _gradient_adtype(adtype), x, constants...
115+
)
116+
hessian_prep = DI.prepare_hessian(target, adtype, x, constants...)
117+
# Buffers pre-allocated from `x`: hot path is zero-allocation on the
118+
# gradient/Hessian outputs, returned arrays alias these slots.
119+
cache = DIHessianCache(
120+
target,
121+
gradient_prep,
122+
hessian_prep,
123+
similar(x),
124+
similar(x, length(x), length(x)),
125+
mode,
126+
)
127+
return Prepared(adtype, evaluator, cache, Val(2))
128+
end
63129
if length(x) == 0
64-
# DI prep crashes on length-0 input (e.g. ForwardDiff `BoundsError`).
65-
# `Val(0)` is an arity sentinel for the `gradient_prep === nothing`
66-
# check below; the AD entry short-circuits before any DI call.
67-
gp, jp = arity === :scalar ? (Val(0), nothing) : (nothing, Val(0))
68-
cache = _wrap_cache(_call_evaluator, gp, jp, Val(length(context)))
130+
cache = if arity === :scalar
131+
DIGradientCache(_call_evaluator, nothing, mode_empty)
132+
else
133+
DIJacobianCache(_call_evaluator, nothing, mode_empty)
134+
end
69135
return Prepared(adtype, evaluator, cache)
70136
end
71137
if arity === :scalar
72138
target, gradient_prep, mode = _prepare_di(DI.prepare_gradient, adtype, x, evaluator)
73-
return Prepared(
74-
adtype, evaluator, _wrap_cache(target, gradient_prep, nothing, mode)
75-
)
139+
return Prepared(adtype, evaluator, DIGradientCache(target, gradient_prep, mode))
76140
end
77141
target, jacobian_prep, mode = _prepare_di(DI.prepare_jacobian, adtype, x, evaluator)
78-
return Prepared(adtype, evaluator, _wrap_cache(target, nothing, jacobian_prep, mode))
142+
return Prepared(adtype, evaluator, DIJacobianCache(target, jacobian_prep, mode))
79143
end
80144

81-
# Hot-path dispatch is by `Mode` (closure vs constants), resolved at compile
82-
# time. The unconstrained method matches every non-`:closure` `Mode` (i.e.
83-
# any `Int N`); `:closure` is strictly more specific and wins for compiled
84-
# tapes. On the constants path we always pass `DI.Constant(eval.f)` plus the
85-
# `N` context constants — `N == 0` collapses the `map` splat to nothing.
86-
@inline _di_value_and_gradient(c::DICache{:closure}, ad, x, _) =
87-
DI.value_and_gradient(c.target, c.gradient_prep, ad, x)
88-
@inline _di_value_and_gradient(c::DICache, ad, x, eval) = DI.value_and_gradient(
145+
# Hot-path dispatch is by cache type + `Mode` (closure vs constants), both
146+
# resolved at compile time. On the constants path we always pass
147+
# `DI.Constant(eval.f)` plus the `N` context constants — `N == 0` collapses
148+
# the `map` splat to nothing.
149+
const _GradientCapable = Union{DIGradientCache,DIHessianCache}
150+
151+
@inline _di_value_and_gradient(
152+
c::Union{DIGradientCache{:closure},DIHessianCache{:closure}}, ad, x, _
153+
) = DI.value_and_gradient(c.target, c.gradient_prep, _gradient_adtype(ad), x)
154+
@inline _di_value_and_gradient(c::_GradientCapable, ad, x, eval) = DI.value_and_gradient(
89155
c.target,
90156
c.gradient_prep,
91-
ad,
157+
_gradient_adtype(ad),
92158
x,
93159
DI.Constant(eval.f),
94160
map(DI.Constant, eval.context)...,
95161
)
96162

97-
@inline _di_value_and_jacobian(c::DICache{:closure}, ad, x, _) =
163+
@inline _di_value_and_jacobian(c::DIJacobianCache{:closure}, ad, x, _) =
98164
DI.value_and_jacobian(c.target, c.jacobian_prep, ad, x)
99-
@inline _di_value_and_jacobian(c::DICache, ad, x, eval) = DI.value_and_jacobian(
165+
@inline _di_value_and_jacobian(c::DIJacobianCache, ad, x, eval) = DI.value_and_jacobian(
100166
c.target,
101167
c.jacobian_prep,
102168
ad,
@@ -105,27 +171,94 @@ end
105171
map(DI.Constant, eval.context)...,
106172
)
107173

174+
@inline _di_value_gradient_and_hessian(c::DIHessianCache{:closure}, ad, x, _) =
175+
DI.value_gradient_and_hessian!(c.target, c.grad_buf, c.hess_buf, c.hessian_prep, ad, x)
176+
@inline _di_value_gradient_and_hessian(c::DIHessianCache, ad, x, eval) =
177+
DI.value_gradient_and_hessian!(
178+
c.target,
179+
c.grad_buf,
180+
c.hess_buf,
181+
c.hessian_prep,
182+
ad,
183+
x,
184+
DI.Constant(eval.f),
185+
map(DI.Constant, eval.context)...,
186+
)
187+
188+
# `value_and_gradient!!`: works on both `DIGradientCache` (order=1 scalar) and
189+
# `DIHessianCache` (order=2). Empty-input caches carry `gradient_prep::Nothing`
190+
# and dispatch to the short-circuit method below; vector-output caches reject.
191+
@inline function AbstractPPL.value_and_gradient!!(
192+
p::Prepared{
193+
<:AbstractADType,
194+
<:VectorEvaluator,
195+
<:Union{DIGradientCache{<:Any,<:Any,Nothing},DIHessianCache{<:Any,<:Any,Nothing}},
196+
},
197+
x::AbstractVector{T},
198+
) where {T<:Real}
199+
Evaluators._check_ad_input(p.evaluator, x)
200+
return (p.evaluator(x), T[])
201+
end
202+
108203
@inline function AbstractPPL.value_and_gradient!!(
109-
p::Prepared{<:AbstractADType,<:VectorEvaluator,<:DICache}, x::AbstractVector{T}
204+
p::Prepared{<:AbstractADType,<:VectorEvaluator,<:_GradientCapable}, x::AbstractVector{T}
110205
) where {T<:Real}
111-
p.cache.gradient_prep === nothing && Evaluators._throw_gradient_needs_scalar()
112206
Evaluators._check_ad_input(p.evaluator, x)
113-
# Bypass DI on length-0 input — DI prep paths fail (e.g. ForwardDiff
114-
# `BoundsError`); typed `T[]` matches the caller's element type.
115-
length(x) == 0 && return (p.evaluator(x), T[])
116207
return _di_value_and_gradient(p.cache, p.adtype, x, p.evaluator)
117208
end
118209

210+
@inline function AbstractPPL.value_and_gradient!!(
211+
::Prepared{<:AbstractADType,<:VectorEvaluator,<:DIJacobianCache},
212+
::AbstractVector{<:Real},
213+
)
214+
return Evaluators._throw_gradient_needs_scalar()
215+
end
216+
119217
@inline function AbstractPPL.value_and_jacobian!!(
120-
p::Prepared{<:AbstractADType,<:VectorEvaluator,<:DICache}, x::AbstractVector{T}
218+
p::Prepared{<:AbstractADType,<:VectorEvaluator,<:DIJacobianCache{<:Any,<:Any,Nothing}},
219+
x::AbstractVector{T},
220+
) where {T<:Real}
221+
Evaluators._check_ad_input(p.evaluator, x)
222+
val = p.evaluator(x)
223+
return (val, similar(x, length(val), 0))
224+
end
225+
226+
@inline function AbstractPPL.value_and_jacobian!!(
227+
p::Prepared{<:AbstractADType,<:VectorEvaluator,<:DIJacobianCache}, x::AbstractVector{T}
121228
) where {T<:Real}
122-
p.cache.jacobian_prep === nothing && Evaluators._throw_jacobian_needs_vector()
123229
Evaluators._check_ad_input(p.evaluator, x)
124-
if length(x) == 0
125-
val = p.evaluator(x)
126-
return (val, similar(x, length(val), 0))
127-
end
128230
return _di_value_and_jacobian(p.cache, p.adtype, x, p.evaluator)
129231
end
130232

233+
@inline function AbstractPPL.value_and_jacobian!!(
234+
::Prepared{<:AbstractADType,<:VectorEvaluator,<:_GradientCapable},
235+
::AbstractVector{<:Real},
236+
)
237+
return Evaluators._throw_jacobian_needs_vector()
238+
end
239+
240+
@inline function AbstractPPL.value_gradient_and_hessian!!(
241+
p::Prepared{
242+
<:AbstractADType,<:VectorEvaluator,<:DIHessianCache{<:Any,<:Any,<:Any,Nothing}
243+
},
244+
x::AbstractVector{T},
245+
) where {T<:Real}
246+
Evaluators._check_ad_input(p.evaluator, x)
247+
return (p.evaluator(x), T[], similar(x, 0, 0))
248+
end
249+
250+
@inline function AbstractPPL.value_gradient_and_hessian!!(
251+
p::Prepared{<:AbstractADType,<:VectorEvaluator,<:DIHessianCache}, x::AbstractVector{T}
252+
) where {T<:Real}
253+
Evaluators._check_ad_input(p.evaluator, x)
254+
return _di_value_gradient_and_hessian(p.cache, p.adtype, x, p.evaluator)
255+
end
256+
257+
@inline function AbstractPPL.value_gradient_and_hessian!!(
258+
::Prepared{<:AbstractADType,<:VectorEvaluator,<:Union{DIGradientCache,DIJacobianCache}},
259+
::AbstractVector{<:Real},
260+
)
261+
return Evaluators._throw_hessian_needs_order_2_prep()
262+
end
263+
131264
end # module

0 commit comments

Comments
 (0)