|
1 | 1 | module AbstractPPLDifferentiationInterfaceExt |
2 | 2 |
|
3 | 3 | using AbstractPPL: AbstractPPL |
4 | | -using AbstractPPL.Evaluators: Evaluators, Prepared, VectorEvaluator |
| 4 | +using AbstractPPL.Evaluators: Evaluators, Prepared, VectorEvaluator, _ad_output_arity |
5 | 5 | using ADTypes: AbstractADType, AutoReverseDiff |
6 | 6 | using DifferentiationInterface: DifferentiationInterface as DI |
7 | 7 |
|
8 | | -# Differentiate only `x`; the evaluator is passed as a `DI.Constant` context so |
9 | | -# that in DynamicPPL the model and other evaluator state stay constant. |
10 | | -@inline _call_evaluator(x, evaluator) = evaluator(x) |
| 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. |
| 13 | +@inline _call_evaluator(x, f::F, ctx::Vararg{Any,N}) where {F,N} = f(x, ctx...) |
11 | 14 |
|
12 | | -struct DICache{F,GP,JP} |
| 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} |
13 | 24 | target::F |
14 | 25 | gradient_prep::GP |
15 | 26 | jacobian_prep::JP |
16 | | - use_context::Bool |
| 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) |
| 29 | + end |
17 | 30 | end |
18 | 31 |
|
19 | 32 | # Compiled ReverseDiff only reuses a compiled tape on the one-argument path; |
20 | 33 | # `DI.Constant` deactivates tape recording, so close the evaluator into the |
21 | | -# target and call DI without contexts. |
| 34 | +# target and call DI without constants. Context (if any) is captured inside |
| 35 | +# the evaluator closure rather than lowered out — the lowered path would also |
| 36 | +# require a closure here, so the wrapper cost is unavoidable for compiled tapes. |
22 | 37 | function _prepare_di(prep::F, adtype::AutoReverseDiff{true}, x, evaluator) where {F} |
23 | 38 | target = Base.Fix2(_call_evaluator, evaluator) |
24 | | - return target, prep(target, adtype, x), false |
| 39 | + return target, prep(target, adtype, x), Val(:closure) |
25 | 40 | end |
26 | 41 |
|
27 | 42 | function _prepare_di(prep::F, adtype::AbstractADType, x, evaluator) where {F} |
28 | | - return _call_evaluator, prep(_call_evaluator, adtype, x, DI.Constant(evaluator)), true |
| 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 | + ) |
29 | 49 | end |
30 | 50 |
|
| 51 | +@inline _wrap_cache(target, gp, jp, ::Val{Mode}) where {Mode} = |
| 52 | + DICache{Mode}(target, gp, jp) |
| 53 | + |
31 | 54 | function AbstractPPL.prepare( |
32 | | - adtype::AbstractADType, problem, x::AbstractVector{<:Real}; check_dims::Bool=true |
| 55 | + adtype::AbstractADType, |
| 56 | + problem, |
| 57 | + x::AbstractVector{<:Real}; |
| 58 | + check_dims::Bool=true, |
| 59 | + context::Tuple=(), |
33 | 60 | ) |
34 | | - evaluator = AbstractPPL.prepare(problem, x; check_dims)::VectorEvaluator |
35 | | - y = evaluator(x) |
36 | | - y isa Union{Number,AbstractVector} || throw( |
37 | | - ArgumentError( |
38 | | - "A prepared AD evaluator must return a scalar or AbstractVector; got $(typeof(y)).", |
39 | | - ), |
40 | | - ) |
| 61 | + evaluator = AbstractPPL.prepare(problem, x; check_dims, context)::VectorEvaluator |
| 62 | + arity = _ad_output_arity(evaluator(x)) |
41 | 63 | if length(x) == 0 |
42 | | - # DI prep crashes on length-0 input (e.g. ForwardDiff `BoundsError`); the |
43 | | - # `Val(0)` sentinel keeps the `gradient_prep === nothing` arity check meaningful. |
44 | | - gp, jp = y isa Number ? (Val(0), nothing) : (nothing, Val(0)) |
45 | | - return Prepared(adtype, evaluator, DICache(_call_evaluator, gp, jp, true)) |
| 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))) |
| 69 | + return Prepared(adtype, evaluator, cache) |
46 | 70 | end |
47 | | - if y isa Number |
48 | | - target, gradient_prep, use_context = _prepare_di( |
49 | | - DI.prepare_gradient, adtype, x, evaluator |
50 | | - ) |
| 71 | + if arity === :scalar |
| 72 | + target, gradient_prep, mode = _prepare_di(DI.prepare_gradient, adtype, x, evaluator) |
51 | 73 | return Prepared( |
52 | | - adtype, evaluator, DICache(target, gradient_prep, nothing, use_context) |
| 74 | + adtype, evaluator, _wrap_cache(target, gradient_prep, nothing, mode) |
53 | 75 | ) |
54 | 76 | end |
55 | | - target, jacobian_prep, use_context = _prepare_di( |
56 | | - DI.prepare_jacobian, adtype, x, evaluator |
57 | | - ) |
58 | | - return Prepared(adtype, evaluator, DICache(target, nothing, jacobian_prep, use_context)) |
| 77 | + target, jacobian_prep, mode = _prepare_di(DI.prepare_jacobian, adtype, x, evaluator) |
| 78 | + return Prepared(adtype, evaluator, _wrap_cache(target, nothing, jacobian_prep, mode)) |
59 | 79 | end |
60 | 80 |
|
| 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( |
| 89 | + c.target, |
| 90 | + c.gradient_prep, |
| 91 | + ad, |
| 92 | + x, |
| 93 | + DI.Constant(eval.f), |
| 94 | + map(DI.Constant, eval.context)..., |
| 95 | +) |
| 96 | + |
| 97 | +@inline _di_value_and_jacobian(c::DICache{:closure}, ad, x, _) = |
| 98 | + 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( |
| 100 | + c.target, |
| 101 | + c.jacobian_prep, |
| 102 | + ad, |
| 103 | + x, |
| 104 | + DI.Constant(eval.f), |
| 105 | + map(DI.Constant, eval.context)..., |
| 106 | +) |
| 107 | + |
61 | 108 | @inline function AbstractPPL.value_and_gradient!!( |
62 | 109 | p::Prepared{<:AbstractADType,<:VectorEvaluator,<:DICache}, x::AbstractVector{T} |
63 | 110 | ) where {T<:Real} |
64 | | - p.cache.gradient_prep === nothing && |
65 | | - throw(ArgumentError("`value_and_gradient!!` requires a scalar-valued function.")) |
66 | | - T <: Integer && Evaluators._reject_integer_input(x) |
67 | | - Evaluators._check_vector_length(p.evaluator.dim, x) |
| 111 | + p.cache.gradient_prep === nothing && Evaluators._throw_gradient_needs_scalar() |
| 112 | + Evaluators._check_ad_input(p.evaluator, x) |
68 | 113 | # Bypass DI on length-0 input — DI prep paths fail (e.g. ForwardDiff |
69 | 114 | # `BoundsError`); typed `T[]` matches the caller's element type. |
70 | 115 | length(x) == 0 && return (p.evaluator(x), T[]) |
71 | | - return if p.cache.use_context |
72 | | - DI.value_and_gradient( |
73 | | - p.cache.target, p.cache.gradient_prep, p.adtype, x, DI.Constant(p.evaluator) |
74 | | - ) |
75 | | - else |
76 | | - DI.value_and_gradient(p.cache.target, p.cache.gradient_prep, p.adtype, x) |
77 | | - end |
| 116 | + return _di_value_and_gradient(p.cache, p.adtype, x, p.evaluator) |
78 | 117 | end |
79 | 118 |
|
80 | 119 | @inline function AbstractPPL.value_and_jacobian!!( |
81 | 120 | p::Prepared{<:AbstractADType,<:VectorEvaluator,<:DICache}, x::AbstractVector{T} |
82 | 121 | ) where {T<:Real} |
83 | | - p.cache.jacobian_prep === nothing && |
84 | | - throw(ArgumentError("`value_and_jacobian!!` requires a vector-valued function.")) |
85 | | - T <: Integer && Evaluators._reject_integer_input(x) |
86 | | - Evaluators._check_vector_length(p.evaluator.dim, x) |
| 122 | + p.cache.jacobian_prep === nothing && Evaluators._throw_jacobian_needs_vector() |
| 123 | + Evaluators._check_ad_input(p.evaluator, x) |
87 | 124 | if length(x) == 0 |
88 | 125 | val = p.evaluator(x) |
89 | 126 | return (val, similar(x, length(val), 0)) |
90 | 127 | end |
91 | | - return if p.cache.use_context |
92 | | - DI.value_and_jacobian( |
93 | | - p.cache.target, p.cache.jacobian_prep, p.adtype, x, DI.Constant(p.evaluator) |
94 | | - ) |
95 | | - else |
96 | | - DI.value_and_jacobian(p.cache.target, p.cache.jacobian_prep, p.adtype, x) |
97 | | - end |
| 128 | + return _di_value_and_jacobian(p.cache, p.adtype, x, p.evaluator) |
98 | 129 | end |
99 | 130 |
|
100 | 131 | end # module |
0 commit comments