diff --git a/HISTORY.md b/HISTORY.md index b42d692c..1ee17908 100644 --- a/HISTORY.md +++ b/HISTORY.md @@ -1,3 +1,9 @@ +## 0.15.4 + +`value_and_gradient!!`, `value_and_jacobian!!`, and `value_gradient_and_hessian!!` now accept a `context=` keyword that overrides, for a single call, the `context` frozen at `prepare` — without re-preparing (#167). The default (`context=nothing`) leaves the hot path unchanged; a `Tuple` override must match the frozen context's element types and shapes — a type mismatch (or a non-`Tuple` override) throws an `ArgumentError` — and reuses the type-keyed cache. Backends that bake context into their prepared state throw an `ArgumentError` for a non-empty-input override (re-`prepare` instead): compiled-tape ReverseDiff (`AutoReverseDiff(; compile=true)`) on all three entry points, and Mooncake's Hessian; Mooncake Jacobian preps are context-free by construction, so only an empty override validates there. Empty input runs no derivative machinery, so no backend rejects an override there (it is still validated). + +Fixed a Mooncake reverse-mode correctness bug ([Mooncake issue #1238](https://github.com/chalk-lab/Mooncake.jl/issues/1238)): when a problem function closed over differentiable data (e.g. a model capturing observed data reached through an in-place solve), a reused prepared cache accumulated that captured data's cotangent and returned an incorrect gradient after the first evaluation. The Mooncake extension now closes the function and its context into a `NoTangent` target, so the captured data carries no cotangent and reuse is correct. + ## 0.15.3 Added the `of` type system: a self-contained, declarative way to specify the shape, element type, and support of model variables. Construct specifications with the exported `of` function or the exported `@of` macro: diff --git a/Project.toml b/Project.toml index f35fb986..d31b454f 100644 --- a/Project.toml +++ b/Project.toml @@ -3,7 +3,7 @@ uuid = "7a57a42e-76ec-4ea3-a279-07e840d6d9cf" keywords = ["probabilistic programming"] license = "MIT" desc = "Common interfaces for probabilistic programming" -version = "0.15.3" +version = "0.15.4" [deps] ADTypes = "47edcb42-4c32-4615-8424-f2b9edc5f35b" diff --git a/docs/src/evaluators.md b/docs/src/evaluators.md index 7e1b1282..f87c5548 100644 --- a/docs/src/evaluators.md +++ b/docs/src/evaluators.md @@ -206,6 +206,32 @@ val, grad = value_and_gradient!!(prepared, [1.0, 2.0, 3.0]) `prepared(x)` evaluates `f(x, context...)`, and `context=()` (the default) preserves the unary `f(x)` shape. +### Overriding the context for a single call + +The context frozen at `prepare` can be replaced for one call by passing a +`context` tuple to `value_and_gradient!!`, `value_and_jacobian!!`, or +`value_gradient_and_hessian!!`, without re-preparing: + +```julia +prepared = prepare(adtype, affine, zeros(3); context=(2.0, 1.0)) +value_and_gradient!!(prepared, x) # uses (2.0, 1.0) +value_and_gradient!!(prepared, x; context=(3.0, 0.0)) # overrides, this call only +``` + +The override must match the prepared context's element types and shapes (the +prepared cache is keyed on types); a type mismatch — or a non-`Tuple` +override — throws an `ArgumentError`. The override is per-call — it does not +mutate the frozen context. Some backends bake the context into their prepared +state and reject an override on a non-empty input with an `ArgumentError` +(re-`prepare` instead): compiled-tape ReverseDiff +(`AutoReverseDiff(; compile=true)`), whose context is baked into its tapes on +all three entry points, and Mooncake's `value_gradient_and_hessian!!`, whose +Hessian cache binds its target by object identity. Mooncake Jacobian preps are +context-free by construction (`prepare` rejects vector-valued problems with +non-empty `context`), so only an empty-`Tuple` override validates there. Empty +input runs no derivative machinery, so no backend rejects an override there — +the override is still validated. + ## Without an AD backend The two-argument form `prepare(problem, x)` is available without any AD diff --git a/ext/AbstractPPLDifferentiationInterfaceExt.jl b/ext/AbstractPPLDifferentiationInterfaceExt.jl index 4ba7c344..9686e962 100644 --- a/ext/AbstractPPLDifferentiationInterfaceExt.jl +++ b/ext/AbstractPPLDifferentiationInterfaceExt.jl @@ -7,10 +7,10 @@ using DifferentiationInterface: DifferentiationInterface as DI # AD target used by every DI cache. `Vararg{Any,N}` with a free `N` forces # specialization on the trailing arity (a bare `Vararg{Any}` would skip it). -# DI invokes this as `_call_evaluator(x, f, c1, …, cN)` on the constants path, -# and as `_call_evaluator(x, evaluator)` (via `Fix2`) on the closure path — +# DI invokes this as `_di_call(x, f, c1, …, cN)` on the constants path, +# and as `_di_call(x, evaluator)` (via `Fix2`) on the closure path — # empty `ctx` then makes the splat a no-op. -@inline _call_evaluator(x, f::F, ctx::Vararg{Any,N}) where {F,N} = f(x, ctx...) +@inline _di_call(x, f::F, ctx::Vararg{Any,N}) where {F,N} = f(x, ctx...) # `Mode` tags the call shape: # * `:closure` — compiled-tape ReverseDiff: target is a `Fix2` closure; the @@ -67,10 +67,10 @@ end # `constants == ()` and the splat at every prep/call site collapses to nothing, # letting prep and call sites share one shape regardless of mode. function _di_call_shape(::AutoReverseDiff{true}, evaluator) - return Base.Fix2(_call_evaluator, evaluator), Val(:closure), () + return Base.Fix2(_di_call, evaluator), Val(:closure), () end function _di_call_shape(::AbstractADType, evaluator) - return _call_evaluator, + return _di_call, Val(length(evaluator.context)), (DI.Constant(evaluator.f), map(DI.Constant, evaluator.context)...) end @@ -100,9 +100,7 @@ function AbstractPPL.prepare( if order == 2 arity === :scalar || Evaluators._throw_hessian_needs_scalar() if length(x) == 0 - cache = DIHessianCache( - _call_evaluator, nothing, nothing, nothing, nothing, mode_empty - ) + cache = DIHessianCache(_di_call, nothing, nothing, nothing, nothing, mode_empty) return Prepared(adtype, evaluator, cache, Val(2)) end # Build both gradient and Hessian preps against the same target so @@ -128,9 +126,9 @@ function AbstractPPL.prepare( end if length(x) == 0 cache = if arity === :scalar - DIGradientCache(_call_evaluator, nothing, mode_empty) + DIGradientCache(_di_call, nothing, mode_empty) else - DIJacobianCache(_call_evaluator, nothing, mode_empty) + DIJacobianCache(_di_call, nothing, mode_empty) end return Prepared(adtype, evaluator, cache) end @@ -148,32 +146,59 @@ end # the `map` splat to nothing. const _GradientCapable = Union{DIGradientCache,DIHessianCache} +# Call-time context override (issue #167) resolves via `Evaluators._resolve_context`. +# Context is threaded as constants whose *values* are read at call time, so an +# override of matching types/shapes just swaps those values — no re-prepare. + +# Compiled-tape ReverseDiff bakes context into its tape, so a call-time override +# can't apply on either closure reject path (gradient and Hessian). +function _throw_compiled_rd_override_unsupported() + throw( + ArgumentError( + "Call-time `context` override is not supported for compiled-tape " * + "ReverseDiff, which bakes the context into its tape. Re-`prepare` " * + "with the new context instead.", + ), + ) +end + @inline _di_value_and_gradient( - c::Union{DIGradientCache{:closure},DIHessianCache{:closure}}, ad, x, _ + c::Union{DIGradientCache{:closure},DIHessianCache{:closure}}, ad, x, _eval, ::Nothing ) = DI.value_and_gradient(c.target, c.gradient_prep, _gradient_adtype(ad), x) -@inline _di_value_and_gradient(c::_GradientCapable, ad, x, eval) = DI.value_and_gradient( - c.target, - c.gradient_prep, - _gradient_adtype(ad), - x, - DI.Constant(eval.f), - map(DI.Constant, eval.context)..., -) +@inline _di_value_and_gradient( + ::Union{DIGradientCache{:closure},DIHessianCache{:closure}}, _ad, _x, _eval, ::Tuple +) = _throw_compiled_rd_override_unsupported() +@inline _di_value_and_gradient(c::_GradientCapable, ad, x, eval, context) = + DI.value_and_gradient( + c.target, + c.gradient_prep, + _gradient_adtype(ad), + x, + DI.Constant(eval.f), + map(DI.Constant, Evaluators._resolve_context(eval, context))..., + ) -@inline _di_value_and_jacobian(c::DIJacobianCache{:closure}, ad, x, _) = +@inline _di_value_and_jacobian(c::DIJacobianCache{:closure}, ad, x, _eval, ::Nothing) = DI.value_and_jacobian(c.target, c.jacobian_prep, ad, x) -@inline _di_value_and_jacobian(c::DIJacobianCache, ad, x, eval) = DI.value_and_jacobian( - c.target, - c.jacobian_prep, - ad, - x, - DI.Constant(eval.f), - map(DI.Constant, eval.context)..., -) +@inline _di_value_and_jacobian(::DIJacobianCache{:closure}, _ad, _x, _eval, ::Tuple) = + _throw_compiled_rd_override_unsupported() +@inline _di_value_and_jacobian(c::DIJacobianCache, ad, x, eval, context) = + DI.value_and_jacobian( + c.target, + c.jacobian_prep, + ad, + x, + DI.Constant(eval.f), + map(DI.Constant, Evaluators._resolve_context(eval, context))..., + ) -@inline _di_value_gradient_and_hessian(c::DIHessianCache{:closure}, ad, x, _) = - DI.value_gradient_and_hessian!(c.target, c.grad_buf, c.hess_buf, c.hessian_prep, ad, x) -@inline _di_value_gradient_and_hessian(c::DIHessianCache, ad, x, eval) = +@inline _di_value_gradient_and_hessian( + c::DIHessianCache{:closure}, ad, x, _eval, ::Nothing +) = DI.value_gradient_and_hessian!(c.target, c.grad_buf, c.hess_buf, c.hessian_prep, ad, x) +@inline _di_value_gradient_and_hessian( + ::DIHessianCache{:closure}, _ad, _x, _eval, ::Tuple +) = _throw_compiled_rd_override_unsupported() +@inline _di_value_gradient_and_hessian(c::DIHessianCache, ad, x, eval, context) = DI.value_gradient_and_hessian!( c.target, c.grad_buf, @@ -182,7 +207,7 @@ const _GradientCapable = Union{DIGradientCache,DIHessianCache} ad, x, DI.Constant(eval.f), - map(DI.Constant, eval.context)..., + map(DI.Constant, Evaluators._resolve_context(eval, context))..., ) # `value_and_gradient!!`: works on both `DIGradientCache` (order=1 scalar) and @@ -194,45 +219,53 @@ const _GradientCapable = Union{DIGradientCache,DIHessianCache} <:VectorEvaluator, <:Union{DIGradientCache{<:Any,<:Any,Nothing},DIHessianCache{<:Any,<:Any,Nothing}}, }, - x::AbstractVector{T}, + x::AbstractVector{T}; + context=nothing, ) where {T<:Real} Evaluators._check_ad_input(p.evaluator, x) - return (p.evaluator(x), T[]) + return (Evaluators._evaluate_with_context(p.evaluator, x, context), T[]) end @inline function AbstractPPL.value_and_gradient!!( - p::Prepared{<:AbstractADType,<:VectorEvaluator,<:_GradientCapable}, x::AbstractVector{T} + p::Prepared{<:AbstractADType,<:VectorEvaluator,<:_GradientCapable}, + x::AbstractVector{T}; + context=nothing, ) where {T<:Real} Evaluators._check_ad_input(p.evaluator, x) - return _di_value_and_gradient(p.cache, p.adtype, x, p.evaluator) + return _di_value_and_gradient(p.cache, p.adtype, x, p.evaluator, context) end @inline function AbstractPPL.value_and_gradient!!( ::Prepared{<:AbstractADType,<:VectorEvaluator,<:DIJacobianCache}, - ::AbstractVector{<:Real}, + ::AbstractVector{<:Real}; + context=nothing, ) return Evaluators._throw_gradient_needs_scalar() end @inline function AbstractPPL.value_and_jacobian!!( p::Prepared{<:AbstractADType,<:VectorEvaluator,<:DIJacobianCache{<:Any,<:Any,Nothing}}, - x::AbstractVector{T}, + x::AbstractVector{T}; + context=nothing, ) where {T<:Real} Evaluators._check_ad_input(p.evaluator, x) - val = p.evaluator(x) + val = Evaluators._evaluate_with_context(p.evaluator, x, context) return (val, similar(x, length(val), 0)) end @inline function AbstractPPL.value_and_jacobian!!( - p::Prepared{<:AbstractADType,<:VectorEvaluator,<:DIJacobianCache}, x::AbstractVector{T} + p::Prepared{<:AbstractADType,<:VectorEvaluator,<:DIJacobianCache}, + x::AbstractVector{T}; + context=nothing, ) where {T<:Real} Evaluators._check_ad_input(p.evaluator, x) - return _di_value_and_jacobian(p.cache, p.adtype, x, p.evaluator) + return _di_value_and_jacobian(p.cache, p.adtype, x, p.evaluator, context) end @inline function AbstractPPL.value_and_jacobian!!( ::Prepared{<:AbstractADType,<:VectorEvaluator,<:_GradientCapable}, - ::AbstractVector{<:Real}, + ::AbstractVector{<:Real}; + context=nothing, ) return Evaluators._throw_jacobian_needs_vector() end @@ -241,22 +274,28 @@ end p::Prepared{ <:AbstractADType,<:VectorEvaluator,<:DIHessianCache{<:Any,<:Any,<:Any,Nothing} }, - x::AbstractVector{T}, + x::AbstractVector{T}; + context=nothing, ) where {T<:Real} Evaluators._check_ad_input(p.evaluator, x) - return (p.evaluator(x), T[], similar(x, 0, 0)) + return ( + Evaluators._evaluate_with_context(p.evaluator, x, context), T[], similar(x, 0, 0) + ) end @inline function AbstractPPL.value_gradient_and_hessian!!( - p::Prepared{<:AbstractADType,<:VectorEvaluator,<:DIHessianCache}, x::AbstractVector{T} + p::Prepared{<:AbstractADType,<:VectorEvaluator,<:DIHessianCache}, + x::AbstractVector{T}; + context=nothing, ) where {T<:Real} Evaluators._check_ad_input(p.evaluator, x) - return _di_value_gradient_and_hessian(p.cache, p.adtype, x, p.evaluator) + return _di_value_gradient_and_hessian(p.cache, p.adtype, x, p.evaluator, context) end @inline function AbstractPPL.value_gradient_and_hessian!!( ::Prepared{<:AbstractADType,<:VectorEvaluator,<:Union{DIGradientCache,DIJacobianCache}}, - ::AbstractVector{<:Real}, + ::AbstractVector{<:Real}; + context=nothing, ) return Evaluators._throw_hessian_needs_order_2_prep() end diff --git a/ext/AbstractPPLForwardDiffExt.jl b/ext/AbstractPPLForwardDiffExt.jl index da8f7c5c..45cbc160 100644 --- a/ext/AbstractPPLForwardDiffExt.jl +++ b/ext/AbstractPPLForwardDiffExt.jl @@ -103,6 +103,22 @@ end # config captured at prep time. @inline _fd_call(x, e::VectorEvaluator) = e.f(x, e.context...) +# `_fd_call` target for a call, over the context resolved by +# `Evaluators._resolve_context` (issue #167): the frozen context (`nothing`) or a +# per-call override. `_fd_call` reads `e.f`/`e.context` directly, so the wrapper's +# `CheckInput` is irrelevant here (`{false}` is fine), and the config's `Tag` — +# keyed on the target type and already bypassed via `Val(false)` — stays valid +# for an override of matching context types. (The empty-input +# primal-under-override lives in `Evaluators._evaluate_with_context`.) +@inline _fd_target(p, context) = Base.Fix2( + _fd_call, + VectorEvaluator{false}( + p.evaluator.f, + p.evaluator.dim, + Evaluators._resolve_context(p.evaluator, context), + ), +) + # `Val(false)` on every hot-path call below skips `ForwardDiff.checktag`. A # user-supplied `adtype.tag` (e.g. DynamicPPL's `DynamicPPLTag` sentinel for # nested AD) has a tag-type parameter that does not equal `typeof(target)`, so @@ -115,19 +131,21 @@ end <:VectorEvaluator, <:Union{FDCache{:scalar,Nothing},FDCache{:hessian,Nothing}}, }, - x::AbstractVector{T}, + x::AbstractVector{T}; + context=nothing, ) where {T<:Real} Evaluators._check_ad_input(p.evaluator, x) - return (p.evaluator(x), T[]) + return (Evaluators._evaluate_with_context(p.evaluator, x, context), T[]) end @inline function AbstractPPL.value_and_gradient!!( p::Prepared{<:AutoForwardDiff,<:VectorEvaluator,<:FDCache{:scalar}}, - x::AbstractVector{<:Real}, + x::AbstractVector{<:Real}; + context=nothing, ) Evaluators._check_ad_input(p.evaluator, x) ForwardDiff.gradient!( - p.cache.result, Base.Fix2(_fd_call, p.evaluator), x, p.cache.config, Val(false) + p.cache.result, _fd_target(p, context), x, p.cache.config, Val(false) ) return (DiffResults.value(p.cache.result), DiffResults.gradient(p.cache.result)) end @@ -136,12 +154,13 @@ end # gradient cache built at prep time — skips the O(n²) Hessian work. @inline function AbstractPPL.value_and_gradient!!( p::Prepared{<:AutoForwardDiff,<:VectorEvaluator,<:FDCache{:hessian}}, - x::AbstractVector{<:Real}, + x::AbstractVector{<:Real}; + context=nothing, ) Evaluators._check_ad_input(p.evaluator, x) ForwardDiff.gradient!( p.cache.gradient_result, - Base.Fix2(_fd_call, p.evaluator), + _fd_target(p, context), x, p.cache.gradient_config, Val(false), @@ -156,7 +175,8 @@ end # the failure mode at compile time. @inline function AbstractPPL.value_and_gradient!!( ::Prepared{<:AutoForwardDiff,<:VectorEvaluator,<:FDCache{:vector}}, - ::AbstractVector{<:Real}, + ::AbstractVector{<:Real}; + context=nothing, ) return Evaluators._throw_gradient_needs_scalar() end @@ -165,27 +185,30 @@ end ::Prepared{ <:AutoForwardDiff,<:VectorEvaluator,<:Union{FDCache{:scalar},FDCache{:hessian}} }, - ::AbstractVector{<:Real}, + ::AbstractVector{<:Real}; + context=nothing, ) return Evaluators._throw_jacobian_needs_vector() end @inline function AbstractPPL.value_and_jacobian!!( p::Prepared{<:AutoForwardDiff,<:VectorEvaluator,<:FDCache{:vector,Nothing}}, - x::AbstractVector{<:Real}, + x::AbstractVector{<:Real}; + context=nothing, ) Evaluators._check_ad_input(p.evaluator, x) - val = p.evaluator(x) + val = Evaluators._evaluate_with_context(p.evaluator, x, context) return (val, similar(x, length(val), 0)) end @inline function AbstractPPL.value_and_jacobian!!( p::Prepared{<:AutoForwardDiff,<:VectorEvaluator,<:FDCache{:vector}}, - x::AbstractVector{<:Real}, + x::AbstractVector{<:Real}; + context=nothing, ) Evaluators._check_ad_input(p.evaluator, x) ForwardDiff.jacobian!( - p.cache.result, Base.Fix2(_fd_call, p.evaluator), x, p.cache.config, Val(false) + p.cache.result, _fd_target(p, context), x, p.cache.config, Val(false) ) return (DiffResults.value(p.cache.result), DiffResults.jacobian(p.cache.result)) end @@ -194,26 +217,31 @@ end ::Prepared{ <:AutoForwardDiff,<:VectorEvaluator,<:Union{FDCache{:scalar},FDCache{:vector}} }, - ::AbstractVector{<:Real}, + ::AbstractVector{<:Real}; + context=nothing, ) return Evaluators._throw_hessian_needs_order_2_prep() end @inline function AbstractPPL.value_gradient_and_hessian!!( p::Prepared{<:AutoForwardDiff,<:VectorEvaluator,<:FDCache{:hessian,Nothing}}, - x::AbstractVector{T}, + x::AbstractVector{T}; + context=nothing, ) where {T<:Real} Evaluators._check_ad_input(p.evaluator, x) - return (p.evaluator(x), T[], similar(x, 0, 0)) + return ( + Evaluators._evaluate_with_context(p.evaluator, x, context), T[], similar(x, 0, 0) + ) end @inline function AbstractPPL.value_gradient_and_hessian!!( p::Prepared{<:AutoForwardDiff,<:VectorEvaluator,<:FDCache{:hessian}}, - x::AbstractVector{<:Real}, + x::AbstractVector{<:Real}; + context=nothing, ) Evaluators._check_ad_input(p.evaluator, x) ForwardDiff.hessian!( - p.cache.result, Base.Fix2(_fd_call, p.evaluator), x, p.cache.config, Val(false) + p.cache.result, _fd_target(p, context), x, p.cache.config, Val(false) ) return ( DiffResults.value(p.cache.result), diff --git a/ext/AbstractPPLMooncakeExt.jl b/ext/AbstractPPLMooncakeExt.jl index 3ff32313..54797148 100644 --- a/ext/AbstractPPLMooncakeExt.jl +++ b/ext/AbstractPPLMooncakeExt.jl @@ -8,26 +8,53 @@ using Mooncake: Mooncake const _MooncakeAD = Union{AutoMooncake,AutoMooncakeForward} -# `NoTangent` stops Mooncake from deriving a `Tangent{...}` for the evaluator -# wrapper's fields on each backward pass. Load-bearing for both: -# * `NamedTupleEvaluator` — passed directly to Mooncake on the NamedTuple -# gradient path. -# * `VectorEvaluator` — wrapped by the order=2 path (Mooncake's Hessian -# API accepts only `AbstractVector` arguments, so -# context is closed over via `VectorEvaluator{false}`). -Mooncake.tangent_type(::Type{<:VectorEvaluator}) = Mooncake.NoTangent +# AD target closing over the problem `f` and its (frozen) `context`, so +# `problem(x, context...)` is presented to Mooncake as a function of `x` alone. +# Its `NoTangent` tangent type keeps `f` and the *entire* `context` tuple +# non-differentiable, which serves both AD modes: +# * every constant argument (`f` and all of `context`, e.g. a model's captured +# observed data) carries no cotangent, so `x` is the only active argument — +# and since the whole context is wrapped in this one `NoTangent` target +# rather than passed to Mooncake as separate arguments, this holds for any +# number of context values (cf. Mooncake issue #1238). An alternative scheme +# that passed `f`/`context` to Mooncake as ordinary arguments without zeroing +# their cotangents would, on a reused reverse-mode cache, accumulate a stale +# cotangent from a differentiable context and corrupt the returned +# `x`-gradient after the first call; `NoTangent` removes that; and +# * the vector-only Hessian API gets a single `AbstractVector` argument. +# Deliberately field-minimal: an extra non-differentiable field (e.g. an `Int`) +# on a `NoTangent` struct defeats Mooncake's forward-mode inference and +# zero-allocation for the wrapper, so the target holds only `f` and `context`. +struct _ADTarget{F,C} + f::F + context::C +end +(t::_ADTarget)(x) = t.f(x, t.context...) + +Mooncake.tangent_type(::Type{<:_ADTarget}) = Mooncake.NoTangent + +# The `_ADTarget` for a call, over the context resolved by +# `Evaluators._resolve_context` (issue #167) — the counterpart of ForwardDiff's +# `_fd_target(p, context)`. On the `nothing` hot path this reconstructs an +# identical immutable struct from the frozen `f`/context (zero-alloc); a `Tuple` +# swaps in the override. The gradient cache re-accepts its target each call, so +# rebuilding it is safe (unlike the identity-bound Hessian cache). +@inline _mc_target(p, context) = + _ADTarget(p.evaluator.f, Evaluators._resolve_context(p.evaluator, context)) + +# `NamedTupleEvaluator` is passed directly to Mooncake on the NamedTuple +# gradient path; the same `NoTangent` reasoning applies to its captured fields. Mooncake.tangent_type(::Type{<:NamedTupleEvaluator}) = Mooncake.NoTangent # Type parameters: # # * `A::Symbol` — `:scalar` / `:vector` for order=1 (output arity), `:hessian` # for order=2. Drives every dispatch decision below. -# * `Target` — `Nothing` for order=1 (Mooncake's gradient/Jacobian API -# takes `f` and `context` as separate args). For `:hessian`, -# a `VectorEvaluator{false}` that closes over `f` and -# `context`, since Mooncake's Hessian API accepts only -# `AbstractVector` arguments. The evaluator's `NoTangent` -# tangent type prevents differentiation of its fields. +# * `Target` — the `_ADTarget` (`f` + `context`) stored only for the +# order=2 Hessian path, whose cache binds it by object +# identity. `Nothing` everywhere else: the order=1 scalar path +# rebuilds its target per call (see `_mc_target`), and +# the vector/Jacobian and empty-input caches have none. # * `C` — the underlying Mooncake cache, or `Nothing` for the # empty-input shortcut. # * `G` — gradient cache populated only at order=2 so the order=1 @@ -45,24 +72,24 @@ struct MooncakeCache{A,Target,C,G} end MooncakeCache{A}(cache) where {A} = MooncakeCache{A}(nothing, cache) -_mooncake_config(adtype) = adtype.config === nothing ? Mooncake.Config() : adtype.config +_mc_config(adtype) = adtype.config === nothing ? Mooncake.Config() : adtype.config # Mooncake exposes separate `prepare_*_cache` entries per AD mode; the call -# shape (callable + active arg + extra args) is the same. Used by the -# NamedTuple path, the order=1 scalar branch, and the order=2 gradient prep. -function _mooncake_gradient_cache(::AutoMooncake, f, x, args...; config) - return Mooncake.prepare_gradient_cache(f, x, args...; config) +# shape (callable + active arg) is the same. Used by the NamedTuple path, the +# order=1 scalar branch, and the order=2 gradient prep. +function _mc_gradient_cache(::AutoMooncake, f, x; config) + return Mooncake.prepare_gradient_cache(f, x; config) end -function _mooncake_gradient_cache(::AutoMooncakeForward, f, x, args...; config) - return Mooncake.prepare_derivative_cache(f, x, args...; config) +function _mc_gradient_cache(::AutoMooncakeForward, f, x; config) + return Mooncake.prepare_derivative_cache(f, x; config) end function AbstractPPL.prepare( adtype::_MooncakeAD, problem, values::NamedTuple; check_dims::Bool=true ) evaluator = AbstractPPL.prepare(problem, values; check_dims)::NamedTupleEvaluator - config = _mooncake_config(adtype) - cache = _mooncake_gradient_cache(adtype, evaluator, values; config) + config = _mc_config(adtype) + cache = _mc_gradient_cache(adtype, evaluator, values; config) return Prepared(adtype, evaluator, cache) end @@ -106,23 +133,17 @@ function AbstractPPL.prepare( ) evaluator = AbstractPPL.prepare(problem, x; check_dims, context)::VectorEvaluator arity = _ad_output_arity(evaluator(x)) - config = _mooncake_config(adtype) + config = _mc_config(adtype) if order == 2 arity === :scalar || Evaluators._throw_hessian_needs_scalar() - # `{false}` skips the per-call shape check — `_check_ad_input` on the - # AD entry already validates `x`. `dim` is unused for `{false}`. - target = VectorEvaluator{false}(evaluator.f, 0, evaluator.context) length(x) == 0 && return Prepared( - adtype, evaluator, MooncakeCache{:hessian}(target, nothing), Val(2) + adtype, evaluator, MooncakeCache{:hessian}(nothing, nothing), Val(2) ) + target = _ADTarget(evaluator.f, evaluator.context) hess_cache = Mooncake.prepare_hessian_cache(target, x; config) # Order=1 gradient cache so `value_and_gradient!!` on the same prep - # skips the Hessian work. Mooncake's `value_and_gradient!!` runs on - # `evaluator.f` with context-as-extra-args, distinct from the wrapped - # `target` used by the Hessian API. - grad_cache = _mooncake_gradient_cache( - adtype, evaluator.f, x, evaluator.context...; config - ) + # skips the Hessian work, built against the same `NoTangent` target. + grad_cache = _mc_gradient_cache(adtype, target, x; config) return Prepared( adtype, evaluator, @@ -145,13 +166,32 @@ function AbstractPPL.prepare( # Compile the tape on the evaluator's `f` and `context` (not the raw # `problem` / `context` kwargs): a downstream override of structural # `prepare` may return a `VectorEvaluator` whose `.f`/`.context` differ - # from the caller-supplied values, and the hot path reads them off the - # evaluator. The reverse-mode vector branch is the only one that can't - # share `_mooncake_gradient_cache` — it needs `prepare_pullback_cache`. - cache = if arity !== :scalar && adtype isa AutoMooncake + # from the caller-supplied values, so the AD target and caches below are + # built from the evaluator's fields. + if arity === :scalar + # Close `f` and `context` into a `NoTangent` target so they carry no + # Mooncake cotangent: `x` is the only active argument. This is what makes + # cache reuse correct for reverse mode — an alternative that passed + # `f`/`context` without zeroing their cotangents would let a reused cache + # accumulate a captured differentiable value's stale cotangent and + # corrupt the returned gradient (issue #1238). Forward mode is unaffected + # but shares the same target; `_ADTarget` (unlike `VectorEvaluator{false}`) + # has no `dim` field, which keeps forward-mode inference and allocations + # intact. + target = _ADTarget(evaluator.f, evaluator.context) + cache = _mc_gradient_cache(adtype, target, x; config) + # The scalar hot path rebuilds the target per call via `_mc_target` + # (cheap, zero-alloc), so the cache need not store it — only the + # identity-bound Hessian cache keeps a `target`. + return Prepared(adtype, evaluator, MooncakeCache{:scalar}(cache)) + end + # Vector arity: `context` is guaranteed empty (checked above), so there is + # nothing to close over and no shadow to carry over. The reverse-mode branch can't + # share `_mc_gradient_cache` — it needs `prepare_pullback_cache`. + cache = if adtype isa AutoMooncake Mooncake.prepare_pullback_cache(evaluator.f, x; config) else - _mooncake_gradient_cache(adtype, evaluator.f, x, evaluator.context...; config) + _mc_gradient_cache(adtype, evaluator.f, x; config) end return Prepared(adtype, evaluator, MooncakeCache{arity}(cache)) end @@ -163,8 +203,14 @@ end # again here would duplicate the second check on every AD call. # (`∂f` is `NoTangent` thanks to the `tangent_type` overload above.) @inline function AbstractPPL.value_and_gradient!!( - p::Prepared{<:_MooncakeAD,<:NamedTupleEvaluator}, values::NamedTuple + p::Prepared{<:_MooncakeAD,<:NamedTupleEvaluator}, values::NamedTuple; context=nothing ) + context === nothing || throw( + ArgumentError( + "Call-time `context` overrides apply only to vector-input " * + "preparations; NamedTuple evaluators take no `context`.", + ), + ) val, (_, grad) = Mooncake.value_and_gradient!!(p.cache, p.evaluator, values) return (val, grad) end @@ -175,42 +221,31 @@ end # `x`. @inline function AbstractPPL.value_and_gradient!!( p::Prepared{<:_MooncakeAD,<:VectorEvaluator,<:MooncakeCache{:scalar,Nothing,Nothing}}, - x::AbstractVector{T}, + x::AbstractVector{T}; + context=nothing, ) where {T<:Real} Evaluators._check_ad_input(p.evaluator, x) - return (p.evaluator(x), T[]) + return (Evaluators._evaluate_with_context(p.evaluator, x, context), T[]) end -# Scalar-gradient hot path. Reverse mode (`Mooncake.Cache`) needs -# `args_to_zero` to mark `x` as the lone active input (`false` on `f`, -# `true` on `x`, `false` on each context value); forward mode -# (`ForwardCache`) derives activity from its seeded argument and rejects -# the kwarg. The `adtype isa AutoMooncake` branch is compile-folded -# since the concrete type lives in `Prepared`'s type parameters. Empty -# `context` collapses the splat and reduces `args_to_zero` to -# `(false, true)`. `tangents[2]` is the `x`-gradient; trailing entries -# (one per context value) are inactive and discarded. -@inline function _mooncake_value_and_gradient(adtype, gcache, e::VectorEvaluator, x) - val, tangents = if adtype isa AutoMooncake - Mooncake.value_and_gradient!!( - gcache, - e.f, - x, - e.context...; - args_to_zero=(false, true, map(_ -> false, e.context)...), - ) - else - Mooncake.value_and_gradient!!(gcache, e.f, x, e.context...) - end +# Scalar-gradient hot path, shared by both AD modes. `target` is the +# `NoTangent` `_ADTarget`, so `x` is the only argument carrying a tangent: +# Mooncake's default zeroing resets `x`, there is nothing else to reset, and the +# context has no shadow that could persist in the reused cotangent buffers. +# Presenting a single active argument also lets forward mode share this path. +# `tangents[2]` is the `x`-gradient (`tangents[1]` is the target's `NoTangent`). +@inline function _mc_value_and_gradient(gcache, target, x) + val, tangents = Mooncake.value_and_gradient!!(gcache, target, x) return (val, tangents[2]) end @inline function AbstractPPL.value_and_gradient!!( p::Prepared{<:_MooncakeAD,<:VectorEvaluator,<:MooncakeCache{:scalar}}, - x::AbstractVector{<:Real}, + x::AbstractVector{<:Real}; + context=nothing, ) Evaluators._check_ad_input(p.evaluator, x) - return _mooncake_value_and_gradient(p.adtype, p.cache.cache, p.evaluator, x) + return _mc_value_and_gradient(p.cache.cache, _mc_target(p, context), x) end # Arity-mismatch errors as dedicated methods so dispatch on @@ -218,7 +253,8 @@ end # time instead of a runtime check on the cache contents. @inline function AbstractPPL.value_and_gradient!!( ::Prepared{<:_MooncakeAD,<:VectorEvaluator,<:MooncakeCache{:vector}}, - ::AbstractVector{<:Real}, + ::AbstractVector{<:Real}; + context=nothing, ) return Evaluators._throw_gradient_needs_scalar() end @@ -229,20 +265,22 @@ end p::Prepared{ <:_MooncakeAD,<:VectorEvaluator,<:MooncakeCache{:hessian,<:Any,Nothing,Nothing} }, - x::AbstractVector{T}, + x::AbstractVector{T}; + context=nothing, ) where {T<:Real} Evaluators._check_ad_input(p.evaluator, x) - return (p.evaluator(x), T[]) + return (Evaluators._evaluate_with_context(p.evaluator, x, context), T[]) end # Order=2 prep also satisfies the order=1 gradient contract via the dedicated # gradient cache built at prep time — skips the O(n²) Hessian work. @inline function AbstractPPL.value_and_gradient!!( p::Prepared{<:_MooncakeAD,<:VectorEvaluator,<:MooncakeCache{:hessian}}, - x::AbstractVector{<:Real}, + x::AbstractVector{<:Real}; + context=nothing, ) Evaluators._check_ad_input(p.evaluator, x) - return _mooncake_value_and_gradient(p.adtype, p.cache.gradient_cache, p.evaluator, x) + return _mc_value_and_gradient(p.cache.gradient_cache, _mc_target(p, context), x) end @inline function AbstractPPL.value_and_jacobian!!( @@ -251,7 +289,8 @@ end <:VectorEvaluator, <:Union{MooncakeCache{:scalar},MooncakeCache{:hessian}}, }, - ::AbstractVector{<:Real}, + ::AbstractVector{<:Real}; + context=nothing, ) return Evaluators._throw_jacobian_needs_vector() end @@ -260,22 +299,25 @@ end # scalar case above; skips Mooncake entirely. @inline function AbstractPPL.value_and_jacobian!!( p::Prepared{<:_MooncakeAD,<:VectorEvaluator,<:MooncakeCache{:vector,Nothing,Nothing}}, - x::AbstractVector{T}, + x::AbstractVector{T}; + context=nothing, ) where {T<:Real} Evaluators._check_ad_input(p.evaluator, x) - val = p.evaluator(x) + val = Evaluators._evaluate_with_context(p.evaluator, x, context) return (val, similar(x, length(val), 0)) end @inline function AbstractPPL.value_and_jacobian!!( p::Prepared{<:_MooncakeAD,<:VectorEvaluator,<:MooncakeCache{:vector}}, - x::AbstractVector{T}, + x::AbstractVector{T}; + context=nothing, ) where {T<:Real} Evaluators._check_ad_input(p.evaluator, x) # Vector arity rejects non-empty `context` at prepare time, so the tape - # is compiled on `problem(x)` and there is no splat or `args_to_zero` to - # propagate. Mooncake's `value_and_jacobian!!` returns `(val, jac)` - # directly with `x` as the only active argument. + # is compiled on `problem(x)` with `x` as the only active argument and the + # frozen context is `()` — the only override that validates is `()` itself. + Evaluators._resolve_context(p.evaluator, context) + # Mooncake's `value_and_jacobian!!` returns `(val, jac)` directly. return Mooncake.value_and_jacobian!!(p.cache.cache, p.evaluator.f, x) end @@ -283,24 +325,41 @@ end # methods below that are strictly more specific, so this catch-all only fires # for `:scalar` / `:vector`. @inline function AbstractPPL.value_gradient_and_hessian!!( - ::Prepared{<:_MooncakeAD,<:VectorEvaluator,<:MooncakeCache}, ::AbstractVector{<:Real} + ::Prepared{<:_MooncakeAD,<:VectorEvaluator,<:MooncakeCache}, + ::AbstractVector{<:Real}; + context=nothing, ) return Evaluators._throw_hessian_needs_order_2_prep() end -# Empty-input shortcut — Mooncake builds no tape for length-zero `x`. +# Empty-input shortcut — Mooncake builds no tape for length-zero `x`, so there +# is no Hessian cache to bind and a `context` override is accepted (it only +# affects the returned value); the identity constraint applies to the tape-built +# path below. @inline function AbstractPPL.value_gradient_and_hessian!!( p::Prepared{<:_MooncakeAD,<:VectorEvaluator,<:MooncakeCache{:hessian,<:Any,Nothing}}, - x::AbstractVector{T}, + x::AbstractVector{T}; + context=nothing, ) where {T<:Real} Evaluators._check_ad_input(p.evaluator, x) - return (p.evaluator(x), T[], similar(x, 0, 0)) + return ( + Evaluators._evaluate_with_context(p.evaluator, x, context), T[], similar(x, 0, 0) + ) end @inline function AbstractPPL.value_gradient_and_hessian!!( p::Prepared{<:_MooncakeAD,<:VectorEvaluator,<:MooncakeCache{:hessian}}, - x::AbstractVector{T}, + x::AbstractVector{T}; + context=nothing, ) where {T<:Real} + context === nothing || throw( + ArgumentError( + "Call-time `context` override is not supported for " * + "`value_gradient_and_hessian!!` with Mooncake, whose Hessian cache " * + "binds its target by object identity. Re-`prepare` with the new " * + "context instead.", + ), + ) Evaluators._check_ad_input(p.evaluator, x) # Mooncake's `value_gradient_and_hessian!!` currently allocates fresh # gradient and Hessian arrays per call despite the `!!` name (unlike its diff --git a/ext/AbstractPPLTestExt.jl b/ext/AbstractPPLTestExt.jl index 05f9cdd1..6e212272 100644 --- a/ext/AbstractPPLTestExt.jl +++ b/ext/AbstractPPLTestExt.jl @@ -5,7 +5,7 @@ using Test: @inferred, @test, @test_broken, @test_throws, @testset """ TestCase(name, tag, f, x_proto; x, value, gradient, jacobian, hessian, - context=(), op, exception, inputs, allocations_safe=true) + context=(), op, exception, inputs, override, allocations_safe=true) Single tagged case for AD conformance testing. The `tag::Symbol` selects how the case is run; the kwargs populate only the fields the tag uses. @@ -17,6 +17,11 @@ Reserved tags (recognised by [`run_testcase`](@ref)): - `:hessian` — order=2 round-trip on scalar output. - `:context` — scalar-output gradient with a non-empty `context::Tuple` passed to `prepare`. + - `:context_override` — the frozen `context` (`gradient`/`hessian` expected) + against a per-call `override::NamedTuple` `(context, value, + gradient, hessian)`. The `gradient_override` / + `hessian_override` [`run_testcase`](@ref) kwargs select + whether each override is honoured or rejected per backend. - `:edge` — error case; `op(prepared, x)` must throw `exception`. - `:cache_reuse` — multiple inputs against a single prepared evaluator (`inputs::Vector{<:NamedTuple}`, with `(x=, value=, @@ -40,6 +45,7 @@ struct TestCase op::Any exception::Any inputs::Any + override::Any allocations_safe::Bool end function TestCase( @@ -56,6 +62,7 @@ function TestCase( op=nothing, exception=nothing, inputs=nothing, + override=nothing, allocations_safe::Bool=true, ) return TestCase( @@ -72,6 +79,7 @@ function TestCase( op, exception, inputs, + override, allocations_safe, ) end @@ -84,6 +92,29 @@ struct VectorValuedProblem end _context_problem(y::AbstractVector{<:Real}, offset) = -0.5 * (y[1] - offset)^2 +# `∂/∂y a·‖y‖² = 2a·y` and its Hessian `2a·I` both depend on the context `a`, so +# a context override is directly observable in the gradient and Hessian. +_affine(y::AbstractVector{<:Real}, a, b) = a * sum(abs2, y) + b + +# Array-valued context: gradient `2w²ᵢyᵢ` and Hessian `Diagonal(2w²)` reflect an +# override of the whole data vector. The empty branch only runs on the +# empty-input shortcuts, which bypass AD entirely. +_weighted(y::AbstractVector{<:Real}, w) = isempty(y) ? zero(eltype(w)) : sum(abs2, y .* w) + +# Vector-output problem for the Jacobian override path: scales the input by the +# first context element, so the Jacobian is `c₁·I` (or `Diagonal(c₁)` for an +# array `c₁`) and an override is directly observable. +_jac_ctx_problem(y::AbstractVector{<:Real}, c1, rest...) = c1 .* y +function _jac_ctx_expected(context, x) + c1 = first(context) + n = length(x) + J = zeros(n, n) + for i in 1:n + J[i, i] = c1 isa AbstractArray ? c1[i] : c1 + end + return J +end + function AbstractPPL.generate_testcases(::Val{:vector}) return ( TestCase( @@ -285,6 +316,46 @@ function AbstractPPL.generate_testcases(::Val{:vector}) ) end +function AbstractPPL.generate_testcases(::Val{:context_override}) + x = [1.0, 2.0, 3.0] + return ( + TestCase( + "affine scalar output, context override", + :context_override, + _affine, + zeros(3); + x=x, + context=(2.0, 1.0), + value=_affine(x, 2.0, 1.0), + gradient=4.0 .* x, # 2a·y with a=2 + hessian=[4.0 0 0; 0 4.0 0; 0 0 4.0], + override=( + context=(3.0, 5.0), + value=_affine(x, 3.0, 5.0), + gradient=6.0 .* x, # 2a·y with a=3 + hessian=[6.0 0 0; 0 6.0 0; 0 0 6.0], + ), + ), + TestCase( + "weighted scalar output, array-context override", + :context_override, + _weighted, + zeros(3); + x=x, + context=([1.0, 2.0, 3.0],), + value=_weighted(x, [1.0, 2.0, 3.0]), + gradient=2.0 .* [1.0, 2.0, 3.0] .^ 2 .* x, + hessian=[2.0 0 0; 0 8.0 0; 0 0 18.0], + override=( + context=([2.0, 1.0, 0.5],), + value=_weighted(x, [2.0, 1.0, 0.5]), + gradient=2.0 .* [2.0, 1.0, 0.5] .^ 2 .* x, + hessian=[8.0 0 0; 0 2.0 0; 0 0 0.5], + ), + ), + ) +end + function AbstractPPL.generate_testcases(::Val{:namedtuple}) return ( TestCase( @@ -388,6 +459,183 @@ function _run( return nothing end +# Context frozen at `prepare` vs a per-call `context=` override, on the +# gradient, Jacobian, and Hessian entry points, plus the invariants that hold on +# every backend. `gradient_override`/`hessian_override` are `:honor` (the value +# is swapped and observed) or `:reject` (the backend bakes context into its +# prepared state and throws for a non-empty input); `jacobian_override` adds +# `:prepare_rejects` for backends that refuse vector arity + non-empty context +# at `prepare` (Mooncake), where only the trivially-matching empty override +# exists. Each backend driver passes the flags that match its caches. Contract +# violations (type/arity mismatch, non-`Tuple`) are validated centrally and +# throw on every backend regardless of the flags. Empty input runs no +# derivative machinery, so no backend rejects an override there — but the +# override is still validated, and integer-eltype inputs are still rejected. +function _run( + ::Val{:context_override}, + case; + adtype, + prepare_fn=AbstractPPL.prepare, + atol=0, + rtol=1e-10, + check_dims::Bool=true, + gradient_override::Symbol=:honor, + hessian_override::Symbol=:honor, + jacobian_override::Symbol=:honor, + kwargs..., +) + ov = case.override + + # --- order=1 gradient --- + prep = prepare_fn(adtype, case.f, case.x_proto; check_dims, context=case.context) + val, grad = AbstractPPL.value_and_gradient!!(prep, case.x) + @test val ≈ case.value atol = atol rtol = rtol + @test grad ≈ case.gradient atol = atol rtol = rtol + if gradient_override === :reject + @test_throws ArgumentError AbstractPPL.value_and_gradient!!( + prep, case.x; context=ov.context + ) + else + val_o, grad_o = AbstractPPL.value_and_gradient!!(prep, case.x; context=ov.context) + @test val_o ≈ ov.value atol = atol rtol = rtol + @test grad_o ≈ ov.gradient atol = atol rtol = rtol + # Matches a prep built with the override context from the start. + fresh = prepare_fn(adtype, case.f, case.x_proto; check_dims, context=ov.context) + @test grad_o ≈ AbstractPPL.value_and_gradient!!(fresh, case.x)[2] atol = atol rtol = + rtol + # The override is per-call: the next default call still uses the frozen context. + @test AbstractPPL.value_and_gradient!!(prep, case.x)[2] ≈ case.gradient atol = atol rtol = + rtol + end + + # --- order=2 gradient+Hessian --- + preph = prepare_fn( + adtype, case.f, case.x_proto; check_dims, context=case.context, order=2 + ) + @test AbstractPPL.value_gradient_and_hessian!!(preph, case.x)[3] ≈ case.hessian atol = + atol rtol = rtol + + # `value_and_gradient!!` on an order=2 prep is a distinct dispatch (the + # gradient cache, not the Hessian one) that also takes the override; it + # follows `gradient_override`, not `hessian_override`. + if gradient_override === :reject + @test_throws ArgumentError AbstractPPL.value_and_gradient!!( + preph, case.x; context=ov.context + ) + else + @test AbstractPPL.value_and_gradient!!(preph, case.x; context=ov.context)[2] ≈ + ov.gradient atol = atol rtol = rtol + end + + if hessian_override === :reject + @test_throws ArgumentError AbstractPPL.value_gradient_and_hessian!!( + preph, case.x; context=ov.context + ) + else + _, grad_h, hess_h = AbstractPPL.value_gradient_and_hessian!!( + preph, case.x; context=ov.context + ) + @test grad_h ≈ ov.gradient atol = atol rtol = rtol + @test hess_h ≈ ov.hessian atol = atol rtol = rtol + # Matches a fresh order=2 prep built with the override context. + freshh = prepare_fn( + adtype, case.f, case.x_proto; check_dims, context=ov.context, order=2 + ) + @test hess_h ≈ AbstractPPL.value_gradient_and_hessian!!(freshh, case.x)[3] atol = + atol rtol = rtol + # Per-call: a subsequent default call still uses the frozen context. + @test AbstractPPL.value_gradient_and_hessian!!(preph, case.x)[3] ≈ case.hessian atol = + atol rtol = rtol + end + + # --- Jacobian override --- + if jacobian_override === :prepare_rejects + # Vector arity + non-empty context is refused at `prepare`, so a + # context-carrying Jacobian prep cannot exist; on a context-free vector + # prep only the trivially-matching empty override validates. + @test_throws ArgumentError prepare_fn( + adtype, _jac_ctx_problem, case.x_proto; check_dims, context=case.context + ) + jprep0 = prepare_fn(adtype, VectorValuedProblem(), case.x_proto; check_dims) + @test AbstractPPL.value_and_jacobian!!(jprep0, case.x; context=())[1] ≈ + VectorValuedProblem()(case.x) atol = atol rtol = rtol + @test_throws ArgumentError AbstractPPL.value_and_jacobian!!( + jprep0, case.x; context=ov.context + ) + else + jprep = prepare_fn( + adtype, _jac_ctx_problem, case.x_proto; check_dims, context=case.context + ) + @test AbstractPPL.value_and_jacobian!!(jprep, case.x)[2] ≈ + _jac_ctx_expected(case.context, case.x) atol = atol rtol = rtol + if jacobian_override === :reject + @test_throws ArgumentError AbstractPPL.value_and_jacobian!!( + jprep, case.x; context=ov.context + ) + else + @test AbstractPPL.value_and_jacobian!!(jprep, case.x; context=ov.context)[2] ≈ + _jac_ctx_expected(ov.context, case.x) atol = atol rtol = rtol + # Per-call: the frozen context is restored afterwards. + @test AbstractPPL.value_and_jacobian!!(jprep, case.x)[2] ≈ + _jac_ctx_expected(case.context, case.x) atol = atol rtol = rtol + end + end + + # --- override contract violations (validated centrally, all backends) --- + # Arity/type mismatches and non-`Tuple` overrides throw an + # `ArgumentError` from `Evaluators._resolve_context` before any backend + # machinery runs (compiled-tape ReverseDiff throws its own `ArgumentError` + # even earlier, so the assertions hold on every backend). + @test_throws ArgumentError AbstractPPL.value_and_gradient!!( + prep, case.x; context=(ov.context..., ov.context[end]) + ) + @test_throws ArgumentError AbstractPPL.value_and_gradient!!( + prep, case.x; context=first(ov.context) + ) + for (i, c) in enumerate(case.context) + c isa AbstractArray || continue + bad_eltype = ntuple( + j -> j == i ? Float32.(case.context[j]) : case.context[j], length(case.context) + ) + @test_throws ArgumentError AbstractPPL.value_and_gradient!!( + prep, case.x; context=bad_eltype + ) + end + + # --- backend-independent invariants --- + # A stray `context=` on a wrong-arity prep surfaces the domain error, not a + # MethodError (reject methods accept and ignore `context`). `VectorValuedProblem` + # is the canonical vector-output (wrong-arity) fixture. + vprep = prepare_fn(adtype, VectorValuedProblem(), case.x_proto; check_dims) + @test_throws ArgumentError AbstractPPL.value_and_gradient!!( + vprep, case.x; context=ov.context + ) + + # Empty input accepts an override on every entry point — no tape/cache built. + e0 = similar(case.x, 0) + empty_val = case.f(e0, ov.context...) + epg = prepare_fn(adtype, case.f, e0; check_dims, context=case.context) + @test AbstractPPL.value_and_gradient!!(epg, e0; context=ov.context)[1] ≈ empty_val atol = + atol rtol = rtol + eph = prepare_fn(adtype, case.f, e0; check_dims, context=case.context, order=2) + @test AbstractPPL.value_gradient_and_hessian!!(eph, e0; context=ov.context)[1] ≈ + empty_val atol = atol rtol = rtol + # A non-`Tuple` override is rejected on the empty-input shortcut too. + @test_throws ArgumentError AbstractPPL.value_and_gradient!!( + epg, e0; context=first(ov.context) + ) + + # Integer-eltype rejection is preserved under an override: the override path + # calls `f` directly, so it must apply the same static guard the evaluator + # callables do. `check_dims=false` isolates the guard (the `{true}` entry + # check already rejects integers before the override is reached). + epg_f = prepare_fn(adtype, case.f, e0; check_dims=false, context=case.context) + @test_throws ArgumentError AbstractPPL.value_and_gradient!!( + epg_f, Int[]; context=ov.context + ) + return nothing +end + function _run( ::Val{:hessian}, case; diff --git a/src/AbstractPPL.jl b/src/AbstractPPL.jl index 3b00f186..f9688fb5 100644 --- a/src/AbstractPPL.jl +++ b/src/AbstractPPL.jl @@ -18,20 +18,26 @@ using .Evaluators: generate_testcases(::Val{group}) Return a tuple of AD conformance test cases for the input-shape `group`. -Reserved groups: `:vector` (vector input) and `:namedtuple` (NamedTuple -input; Mooncake-only). Iterate and pass each to [`run_testcase`](@ref). +Reserved groups: `:vector` (vector input), `:namedtuple` (NamedTuple +input; Mooncake-only), and `:context_override` (call-time `context` +overrides). Iterate and pass each to [`run_testcase`](@ref). Implemented by the `Test` extension (`AbstractPPLTestExt`). """ function generate_testcases end """ run_testcase(case; adtype, prepare_fn=AbstractPPL.prepare, atol=0, rtol=1e-10, - check_dims=true, type_stability=:skip, allocations=:skip) + check_dims=true, type_stability=:skip, allocations=:skip, + gradient_override=:honor, hessian_override=:honor, + jacobian_override=:honor) Run a single conformance case against an AD backend. `type_stability` and `allocations` accept `:skip` / `:test` / `:broken` — `:test` asserts the invariant, `:broken` marks it `@test_broken` (use for backends with known -regressions). Implemented by the `Test` extension. +regressions). The `*_override` kwargs apply to the `:context_override` group +and declare whether the backend honours or rejects a call-time `context` +override on each entry point (see the `Test` extension's `TestCase` docs). +Implemented by the `Test` extension. """ function run_testcase end diff --git a/src/evaluators/Evaluators.jl b/src/evaluators/Evaluators.jl index 124c7134..b88b09b5 100644 --- a/src/evaluators/Evaluators.jl +++ b/src/evaluators/Evaluators.jl @@ -113,32 +113,61 @@ function prepare( end """ - value_and_gradient!!(prepared, x::AbstractVector{<:Real}) + value_and_gradient!!(prepared, x::AbstractVector{<:Real}; context=nothing) Return `(value, gradient)` for a scalar-valued evaluator, potentially reusing internal cache buffers of `prepared`. The returned gradient may alias `prepared`'s internal storage; copy if you need to retain it past the next call. + +By default the `context` frozen at `prepare` is used. Pass a `Tuple` as +`context` to override it for a single call without re-preparing — it must match +the prepared context's element types and shapes, since the prepared cache is +keyed on types. A type mismatch, or a non-`Tuple` override, throws an +`ArgumentError`. The override is per-call and does not mutate the frozen +context. Overrides apply to vector-input preparations; a NamedTuple prep +accepts only `context=nothing`. +Compiled-tape ReverseDiff (`AutoReverseDiff(; compile=true)`) bakes the context +into its tape and throws if an override is supplied for a non-empty input. +Empty input (`length(x) == 0`) runs no derivative machinery — the value is +computed directly — so no backend rejects an override there; the override is +still validated against the frozen context. """ function value_and_gradient!! end """ - value_and_jacobian!!(prepared, x::AbstractVector{<:Real}) + value_and_jacobian!!(prepared, x::AbstractVector{<:Real}; context=nothing) Return `(value::AbstractVector, jacobian::AbstractMatrix)` for a vector-valued evaluator, potentially reusing internal cache buffers. The returned arrays may alias `prepared`'s internal storage; copy if needed. The Jacobian has shape `(length(value), length(x))`. + +`context` may override the context frozen at `prepare` for a single call, under +the same contract as [`value_and_gradient!!`](@ref). Compiled-tape ReverseDiff +(`AutoReverseDiff(; compile=true)`) bakes the context into its tape and throws +for a non-empty-input override. Mooncake Jacobian preps are context-free by +construction (`prepare` rejects vector-valued problems with non-empty +`context`), so only an empty-`Tuple` override validates there. """ function value_and_jacobian!! end """ - value_gradient_and_hessian!!(prepared, x::AbstractVector{<:Real}) + value_gradient_and_hessian!!(prepared, x::AbstractVector{<:Real}; context=nothing) Return `(value, gradient::AbstractVector, hessian::AbstractMatrix)` for a scalar-valued evaluator prepared with `order=2`, potentially reusing internal cache buffers. The returned gradient and Hessian may alias `prepared`'s internal storage; copy if you need to retain them past the next call. The Hessian has shape `(length(x), length(x))`. + +`context` may override the context frozen at `prepare` for a single call, under +the same contract as [`value_and_gradient!!`](@ref), except that on a non-empty +input two backends throw for the Hessian and require re-`prepare` instead: +compiled-tape ReverseDiff (`AutoReverseDiff(; compile=true)`), which bakes the +context into its tape, and Mooncake, whose Hessian cache binds its target by +object identity. As for the gradient, empty input runs no machinery, so no +backend rejects an override there (it is still validated against the frozen +context). """ function value_gradient_and_hessian!! end @@ -247,6 +276,58 @@ function _check_ad_input(e::VectorEvaluator{true}, x::AbstractVector{T}) where { end _check_ad_input(::VectorEvaluator{false}, ::AbstractVector) = nothing +# Primal value under a possibly-overridden call-time context, shared by the +# AD-backend extensions' empty-input shortcuts (which bypass the backend). +# `nothing` keeps the frozen context via the evaluator (honouring its +# `CheckInput`); a `Tuple` is validated by `_resolve_context` and applied to +# `f` directly, with the same static integer-eltype rejection as the evaluator +# callables (on the `{false}` path there is otherwise no check between here +# and `f`); anything else is rejected with the same error as the AD paths. +@inline _evaluate_with_context(e::VectorEvaluator, x, ::Nothing) = e(x) +@inline function _evaluate_with_context( + e::VectorEvaluator, x::AbstractVector{T}, context::Tuple +) where {T} + T <: Integer && _reject_integer_input(x) + return e.f(x, _resolve_context(e, context)...) +end +function _evaluate_with_context(::VectorEvaluator, _, context) + throw( + ArgumentError( + "A call-time `context` override must be a `Tuple`; got $(typeof(context))." + ), + ) +end + +# Resolve a call-time `context` override into the context tuple a backend builds +# its AD target from: `nothing` (the default) keeps the context frozen at +# `prepare`; a `Tuple` replaces it for this call (issue #167). Shared by the +# AD-backend override paths so the `nothing`/`Tuple` dispatch and the override +# validation live in one place; each backend wraps the result in its own target +# (`Constant`s, `_ADTarget`, a `Fix2` evaluator). +# +# An override must match the frozen context's element types, since the prepared +# caches are keyed on types: `typeof === ` catches arity and every element type +# in one comparison that constant-folds away when the types match. Array sizes +# are not checked — keeping them fixed is the caller's responsibility. +@inline _resolve_context(e::VectorEvaluator, ::Nothing) = e.context +@inline function _resolve_context(e::VectorEvaluator, context::Tuple) + typeof(context) === typeof(e.context) || throw( + ArgumentError( + "Call-time `context` override does not match the context frozen at " * + "`prepare`: expected `$(typeof(e.context))`, got `$(typeof(context))`. " * + "The prepared caches are keyed on these types; re-`prepare` to change them.", + ), + ) + return context +end +function _resolve_context(::VectorEvaluator, context) + throw( + ArgumentError( + "A call-time `context` override must be a `Tuple`; got $(typeof(context))." + ), + ) +end + # Both bodies rely on `T <: Integer` being a static check so the AD hot path # (Float/dual `T`) elides the branch; the `{false}` callable additionally skips # `_check_vector_length` since AD libraries pass length-matching dual inputs. diff --git a/test/ext/differentiationinterface/main.jl b/test/ext/differentiationinterface/main.jl index e709bbd9..d61dbf5a 100644 --- a/test/ext/differentiationinterface/main.jl +++ b/test/ext/differentiationinterface/main.jl @@ -99,4 +99,24 @@ quadratic(x::AbstractVector{<:Real}) = sum(xi -> xi^2, x) @test grad_ctx ≈ [4.0, 8.0, 12.0] @test hess_ctx ≈ [4.0 0 0; 0 4.0 0; 0 0 4.0] end + + # Non-compiled ReverseDiff threads context as call-time `Constant`s and + # honours an override on all three entry points; compiled-tape ReverseDiff + # bakes context into its tapes and rejects a non-empty override on each of + # them. Empty input is accepted on both (the runner checks it regardless of + # the flags). + @testset "call-time context override (#167)" begin + for case in generate_testcases(Val(:context_override)) + run_testcase(case; adtype=AutoReverseDiff(), atol=1e-6, rtol=1e-6) + run_testcase( + case; + adtype=AutoReverseDiff(; compile=true), + atol=1e-6, + rtol=1e-6, + gradient_override=:reject, + hessian_override=:reject, + jacobian_override=:reject, + ) + end + end end diff --git a/test/ext/forwarddiff/main.jl b/test/ext/forwarddiff/main.jl index 5173339b..4eac29ab 100644 --- a/test/ext/forwarddiff/main.jl +++ b/test/ext/forwarddiff/main.jl @@ -47,4 +47,12 @@ using Test @test val ≈ 5.0 @test grad ≈ [2.0, 4.0] end + + # ForwardDiff rebuilds the target per call, so both the gradient and Hessian + # entry points honour a call-time `context` override. + @testset "call-time context override (#167)" begin + for case in generate_testcases(Val(:context_override)) + run_testcase(case; adtype=AutoForwardDiff(), atol=1e-6, rtol=1e-6) + end + end end diff --git a/test/ext/mooncake/Project.toml b/test/ext/mooncake/Project.toml index 6a5c2039..a44d5f69 100644 --- a/test/ext/mooncake/Project.toml +++ b/test/ext/mooncake/Project.toml @@ -1,6 +1,7 @@ [deps] AbstractPPL = "7a57a42e-76ec-4ea3-a279-07e840d6d9cf" ADTypes = "47edcb42-4c32-4615-8424-f2b9edc5f35b" +LinearAlgebra = "37e2e46d-f89d-539d-b4ee-838fcccc9c8e" Mooncake = "da2b9cff-9c12-43a0-ae48-6db2b0edb7d6" Pkg = "44cfe95a-1eb2-52ea-b672-e2afdf69b78f" Test = "8dfed614-e22c-5e08-85e1-65c5234f0b40" diff --git a/test/ext/mooncake/main.jl b/test/ext/mooncake/main.jl index 50c2f2d3..b85ce6fc 100644 --- a/test/ext/mooncake/main.jl +++ b/test/ext/mooncake/main.jl @@ -6,6 +6,7 @@ Pkg.instantiate() using AbstractPPL: AbstractPPL, prepare, generate_testcases, run_testcase, value_and_gradient!! using ADTypes: AutoMooncake, AutoMooncakeForward +using LinearAlgebra: LowerTriangular, I using Mooncake using Test @@ -13,8 +14,8 @@ using Test # * `value_and_jacobian!!` allocates fresh cotangent/Jacobian buffers on # every call (both modes); forward-mode Jacobian return type infers as # `Tuple{Any, Union{Array{T,3}, Matrix}}`. -# * `value_and_gradient!!` on a forward-mode context-lowered prep splats -# `args_to_zero` per call and allocates; forward mode also fails inference. +# * `value_and_gradient!!` on a forward-mode context-lowered prep allocates +# per call and also fails inference. function _mooncake_alloc(case, adtype) if case.tag === :vector && case.jacobian !== nothing return :broken @@ -59,6 +60,13 @@ end for case in generate_testcases(Val(:namedtuple)) run_testcase(case; adtype, atol=1e-6, rtol=1e-6) end + # NamedTuple preps take no `context`: `nothing` is accepted, a + # `Tuple` override is rejected. + nt_prep = prepare(adtype, v -> v.a^2, (a=0.5,)) + @test value_and_gradient!!(nt_prep, (a=0.5,); context=nothing)[1] ≈ 0.25 + @test_throws ArgumentError value_and_gradient!!( + nt_prep, (a=0.5,); context=(1.0,) + ) end end @@ -132,4 +140,67 @@ end @test_throws r"dense vector" prepare(AutoMooncake(), problem, v) @test_throws r"dense vector" prepare(AutoMooncakeForward(), problem, v) end + + # Reused reverse-mode cache with a differentiable `context` routed through an + # in-place linear solve — the shape that reproduces Mooncake issue #1238. An + # extension that passes `context` to Mooncake as an ordinary, non-zeroed + # argument leaves the data's cotangent to accumulate in a reused cache; the + # `\` reverse rule reads that cotangent back, so the *returned* `x`-gradient + # is corrupted after the first call. + # Closing `f`/`context` into a `NoTangent` `_ADTarget` makes the context + # non-differentiable, fixing it. The in-place `\` is essential — hand-unrolled + # forward substitution never hits the leaking reverse rule (see PR #177). The + # leak is reverse-mode only, so the `AutoMooncake` subtest carries the teeth — + # its reused-vs-fresh oracle fails against the pre-fix scheme and passes under + # the `_ADTarget` fix; the `AutoMooncakeForward` subtest never leaked and + # guards mode parity. + @testset "reused cache with a differentiable context does not leak (#1238)" begin + M(p) = LowerTriangular([p[1] 0 0; p[2] p[1] 0; p[3] p[2] p[1]] + I) + solve_problem(x, data) = sum(abs2, M(x) \ data) + data = [1.0, 2.0, 3.0] + xA = [0.3, 0.5, 0.2] + xB = [1.5, 2.0, -1.0] + # Ground-truth `x`-gradient at `xB` by central differences. + fd = map(eachindex(xB)) do i + e = zeros(length(xB)) + e[i] = 1.0e-6 + (solve_problem(xB .+ e, data) - solve_problem(xB .- e, data)) / 2.0e-6 + end + @testset "$ad" for ad in ( + AutoMooncake(; config=nothing), AutoMooncakeForward(; config=nothing) + ) + # A fresh prep's first evaluation is always correct — the oracle. + fresh = prepare(ad, solve_problem, xB; check_dims=false, context=(data,)) + _, g_fresh = value_and_gradient!!(fresh, xB) + @test g_fresh ≈ fd rtol = 1.0e-5 + # Reuse one prep across two inputs; the second call is the one a + # pre-fix cache corrupted. + prep = prepare(ad, solve_problem, xA; check_dims=false, context=(data,)) + value_and_gradient!!(prep, xA) + _, g_reused = value_and_gradient!!(prep, xB) + @test g_reused ≈ g_fresh + @test g_reused ≈ fd rtol = 1.0e-5 + end + end + + # Mooncake honours a call-time `context` override on the gradient path (the + # target is rebuilt per call) but rejects it on the Hessian path (whose cache + # binds its target by object identity). Both AD modes share the gradient + # path. `check_dims=false` matches the other Mooncake cases. + @testset "call-time context override (#167)" begin + @testset "$ad" for ad in ( + AutoMooncake(; config=nothing), AutoMooncakeForward(; config=nothing) + ) + for case in generate_testcases(Val(:context_override)) + run_testcase( + case; + adtype=ad, + check_dims=false, + rtol=1.0e-6, + hessian_override=:reject, + jacobian_override=:prepare_rejects, + ) + end + end + end end