From 792dc5a5a37f6e63c20d4ae297ca5d7da2e1d60a Mon Sep 17 00:00:00 2001 From: Hong Ge Date: Tue, 7 Jul 2026 17:48:27 +0100 Subject: [PATCH 1/3] Harden reused reverse-mode Mooncake caches; unify the two AD modes MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit The order=1 scalar-gradient path passed the problem function `f` and any `context` to Mooncake as ordinary arguments, without zeroing their cotangents. When `context` carries *differentiable* data reached through an in-place linear solve — whose reverse rule reads its argument's cotangent back — a reused reverse-mode cache let that cotangent accumulate across calls and returned an incorrect `x`-gradient after the first evaluation (Mooncake issue #1238); only the first evaluation on a reused prep was correct. Close `f` and `context` into a `NoTangent`-typed `_ADTarget` instead, so they carry no cotangent by construction and `x` is the only active argument — nothing else for Mooncake to zero, and no leftover cotangent to accumulate. This also unifies the `AutoMooncake` / `AutoMooncakeForward` code paths on that target, dropping the per-mode branching. The target is deliberately field-minimal (`f` + `context`, no `dim`): an extra non-differentiable field defeats Mooncake's forward-mode type inference and zero-allocation. A regression test (a differentiable `context` routed through a triangular solve on a reused prep, checked against a fresh prep and finite differences) guards it. Co-Authored-By: Claude Opus 4.8 (1M context) --- ext/AbstractPPLMooncakeExt.jl | 142 ++++++++++++++++++--------------- test/ext/mooncake/Project.toml | 1 + test/ext/mooncake/main.jl | 47 ++++++++++- 3 files changed, 123 insertions(+), 67 deletions(-) diff --git a/ext/AbstractPPLMooncakeExt.jl b/ext/AbstractPPLMooncakeExt.jl index 3ff32313..8bbd9004 100644 --- a/ext/AbstractPPLMooncakeExt.jl +++ b/ext/AbstractPPLMooncakeExt.jl @@ -8,26 +8,42 @@ 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 + +# `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` closing over `f` and `context`, shared by +# the order=1 scalar-gradient path (both AD modes) and the +# order=2 Hessian path. `Nothing` for the vector/Jacobian and +# empty-input caches. # * `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 +61,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 +122,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 +155,29 @@ 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) + return Prepared(adtype, evaluator, MooncakeCache{:scalar}(target, 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 @@ -181,27 +207,14 @@ end return (p.evaluator(x), 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 @@ -210,7 +223,7 @@ end x::AbstractVector{<:Real}, ) 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, p.cache.target, x) end # Arity-mismatch errors as dedicated methods so dispatch on @@ -242,7 +255,7 @@ end x::AbstractVector{<:Real}, ) 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, p.cache.target, x) end @inline function AbstractPPL.value_and_jacobian!!( @@ -273,9 +286,8 @@ end ) 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. + # Mooncake's `value_and_jacobian!!` returns `(val, jac)` directly. return Mooncake.value_and_jacobian!!(p.cache.cache, p.evaluator.f, x) 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..505eeb9b 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 @@ -132,4 +133,46 @@ 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 end From 74b548820bacf4f8c92c2d065710757cd64db1a8 Mon Sep 17 00:00:00 2001 From: Hong Ge Date: Tue, 7 Jul 2026 17:48:27 +0100 Subject: [PATCH 2/3] Support call-time `context` overrides in the AD entry points (#167) MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit `context` was frozen at `prepare`. Allow `value_and_gradient!!` and `value_gradient_and_hessian!!` to take a `context=` keyword that overrides it for a single call without re-preparing, across the Mooncake, DifferentiationInterface, and ForwardDiff extensions. `context=nothing` (the default) reuses the target frozen at `prepare` — the zero-allocation, type-stable hot path is unchanged. A `Tuple` of matching element types and shapes rebuilds the target (or swaps the prepared constants' values) and reuses the type-keyed cache. The override is per-call and does not mutate the frozen context. The `nothing`/`Tuple` resolution is unified into one `Evaluators._resolve_context`; each backend builds its own target from the result (constants, an `_ADTarget`, or a `Fix2` evaluator). Two backends cannot reuse the prep for an override and throw a clear `ArgumentError` (re-`prepare` instead): compiled-tape ReverseDiff, whose context is baked into the tape, and Mooncake's Hessian, whose cache binds its target by object identity. Empty input runs no derivative machinery, so an override is always accepted there. Consolidate the AD conformance tests: the per-backend override testsets become a shared `:context_override` TestCase tag + runner (with `gradient_override`/`hessian_override` flags), and the differentiable-context reused-cache regression test becomes a shared `:cache_reuse_context` case run across all backends. Tested across all three backends, including the rejection and empty-input paths. Co-Authored-By: Claude Opus 4.8 (1M context) --- HISTORY.md | 6 + Project.toml | 2 +- docs/src/evaluators.md | 21 +++ ext/AbstractPPLDifferentiationInterfaceExt.jl | 98 ++++++++---- ext/AbstractPPLForwardDiffExt.jl | 49 ++++-- ext/AbstractPPLMooncakeExt.jl | 82 +++++++--- ext/AbstractPPLTestExt.jl | 141 +++++++++++++++++- src/evaluators/Evaluators.jl | 39 ++++- test/ext/differentiationinterface/main.jl | 18 +++ test/ext/forwarddiff/main.jl | 8 + test/ext/mooncake/main.jl | 16 ++ 11 files changed, 408 insertions(+), 72 deletions(-) diff --git a/HISTORY.md b/HISTORY.md index b42d692c..608eee94 100644 --- a/HISTORY.md +++ b/HISTORY.md @@ -1,3 +1,9 @@ +## 0.15.4 + +`value_and_gradient!!` and `value_gradient_and_hessian!!` now accept a `context=` keyword that overrides, for a single call, the `context` frozen at `prepare` — without re-preparing. The default (`context=nothing`) leaves the zero-allocation hot path unchanged; a `Tuple` of matching element types and shapes reuses the type-keyed cache. Two backends bake the context into their prepared state and throw an `ArgumentError` for a non-empty-input override (re-`prepare` instead): compiled-tape ReverseDiff (`AutoReverseDiff(; compile=true)`) and Mooncake's Hessian. Empty input runs no derivative machinery, so an override is always accepted there. + +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..30d119fd 100644 --- a/docs/src/evaluators.md +++ b/docs/src/evaluators.md @@ -206,6 +206,27 @@ 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!!` 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), and it is per-call — it does not mutate the +frozen context. Two 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 the tape, and Mooncake's `value_gradient_and_hessian!!`, +whose Hessian cache binds its target by object identity. Empty input runs no +derivative machinery, so an override is always accepted there. + ## 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..35768f90 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,17 +146,37 @@ 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, _) = DI.value_and_jacobian(c.target, c.jacobian_prep, ad, x) @@ -171,9 +189,13 @@ const _GradientCapable = Union{DIGradientCache,DIHessianCache} map(DI.Constant, 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 +204,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,22 +216,26 @@ 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 @@ -241,22 +267,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..0e532064 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/shapes. (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 @@ -194,26 +214,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 8bbd9004..dade72f0 100644 --- a/ext/AbstractPPLMooncakeExt.jl +++ b/ext/AbstractPPLMooncakeExt.jl @@ -30,8 +30,18 @@ struct _ADTarget{F,C} 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 @@ -40,10 +50,11 @@ Mooncake.tangent_type(::Type{<:NamedTupleEvaluator}) = Mooncake.NoTangent # # * `A::Symbol` — `:scalar` / `:vector` for order=1 (output arity), `:hessian` # for order=2. Drives every dispatch decision below. -# * `Target` — the `_ADTarget` closing over `f` and `context`, shared by -# the order=1 scalar-gradient path (both AD modes) and the -# order=2 Hessian path. `Nothing` for the vector/Jacobian and -# empty-input caches. +# * `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 @@ -163,13 +174,16 @@ function AbstractPPL.prepare( # 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. + # 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) - return Prepared(adtype, evaluator, MooncakeCache{:scalar}(target, cache)) + # 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 @@ -201,10 +215,11 @@ 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, shared by both AD modes. `target` is the @@ -220,10 +235,11 @@ 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 _mc_value_and_gradient(p.cache.cache, p.cache.target, x) + return _mc_value_and_gradient(p.cache.cache, _mc_target(p, context), x) end # Arity-mismatch errors as dedicated methods so dispatch on @@ -231,7 +247,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 @@ -242,20 +259,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 _mc_value_and_gradient(p.cache.gradient_cache, p.cache.target, x) + return _mc_value_and_gradient(p.cache.gradient_cache, _mc_target(p, context), x) end @inline function AbstractPPL.value_and_jacobian!!( @@ -295,24 +314,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..7923f3fa 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,10 @@ 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 + function AbstractPPL.generate_testcases(::Val{:vector}) return ( TestCase( @@ -285,6 +297,29 @@ 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], + ), + ), + ) +end + function AbstractPPL.generate_testcases(::Val{:namedtuple}) return ( TestCase( @@ -388,6 +423,110 @@ function _run( return nothing end +# Context frozen at `prepare` vs a per-call `context=` override, on both the +# gradient 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); each backend driver passes +# the pair that matches its caches. Empty input runs no derivative machinery, so +# an override is always accepted there regardless of the flags. +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, + 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 + + # --- 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 + return nothing +end + function _run( ::Val{:hessian}, case; diff --git a/src/evaluators/Evaluators.jl b/src/evaluators/Evaluators.jl index 124c7134..2f71da5b 100644 --- a/src/evaluators/Evaluators.jl +++ b/src/evaluators/Evaluators.jl @@ -113,11 +113,20 @@ 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. The override is per-call and does not mutate the frozen context. +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 a `context` override is always accepted there. """ function value_and_gradient!! end @@ -132,13 +141,21 @@ The Jacobian has shape `(length(value), length(x))`. 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 an +override is always accepted there. """ function value_gradient_and_hessian!! end @@ -247,6 +264,24 @@ 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` calls `f` with the override directly (the per-call +# shape check having already run via `_check_ad_input` at the entry point). +@inline _evaluate_with_context(e::VectorEvaluator, x, ::Nothing) = e(x) +@inline _evaluate_with_context(e::VectorEvaluator, x, context::Tuple) = e.f(x, context...) + +# 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 lives 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 +# and shapes, since the prepared cache is keyed on types. +@inline _resolve_context(e::VectorEvaluator, ::Nothing) = e.context +@inline _resolve_context(::VectorEvaluator, context::Tuple) = context + # 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..bb811809 100644 --- a/test/ext/differentiationinterface/main.jl +++ b/test/ext/differentiationinterface/main.jl @@ -99,4 +99,22 @@ 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 both entry points; compiled-tape ReverseDiff bakes + # context into its tape and rejects a non-empty override. 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, + ) + 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/main.jl b/test/ext/mooncake/main.jl index 505eeb9b..dc306297 100644 --- a/test/ext/mooncake/main.jl +++ b/test/ext/mooncake/main.jl @@ -175,4 +175,20 @@ end @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 + ) + end + end + end end From 74aa79519b0549414f52486de18e77d3ff7601d4 Mon Sep 17 00:00:00 2001 From: Xianda Sun Date: Fri, 10 Jul 2026 08:08:14 +0100 Subject: [PATCH 3/3] Validate context overrides centrally; extend them to `value_and_jacobian!!` MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit The override contract stated in the docs — a call-time `context` must match the frozen context's element types and shapes — was enforced by no one: `_resolve_context` was a pass-through, and violations behaved differently on every backend (silently honoured on ForwardDiff, an internals-leaking `PreparedCacheSpecError` on Mooncake, a raw `MethodError` for non-`Tuple` input). `_resolve_context` now validates every override in one place: `typeof === ` catches arity and element types, top-level array elements are compared by `size`, and a non-`Tuple` override gets a dedicated `ArgumentError`. The empty-input path (`_evaluate_with_context`) routes through the same validation and now applies the evaluator callables' static integer-eltype rejection, which the raw `e.f` call used to bypass. `value_and_jacobian!!` gains the `context=` keyword, completing #167 across all three entry points: DI threads the resolved override as `Constant`s (compiled-tape ReverseDiff rejects, as on the other entry points), ForwardDiff rebuilds its target via the existing `_fd_target`, and Mooncake — whose vector-arity preps are context-free by construction — accepts exactly the empty override via the central check, with no special-casing. Tests: an array-valued-context override fixture (previously untested), a Jacobian override fixture with a `jacobian_override` (:honor/:reject/:prepare_rejects) runner flag, contract-violation cases (arity/eltype/shape/non-`Tuple`) asserted on every backend, and a regression for the integer-guard bypass. HISTORY also drops the "zero-allocation hot path" claim, which held only for scalar contexts. --- HISTORY.md | 2 +- docs/src/evaluators.md | 22 +-- ext/AbstractPPLDifferentiationInterfaceExt.jl | 35 +++-- ext/AbstractPPLForwardDiffExt.jl | 13 +- ext/AbstractPPLMooncakeExt.jl | 15 ++- ext/AbstractPPLTestExt.jl | 126 +++++++++++++++++- src/evaluators/Evaluators.jl | 94 +++++++++++-- test/ext/differentiationinterface/main.jl | 8 +- test/ext/mooncake/main.jl | 7 +- 9 files changed, 265 insertions(+), 57 deletions(-) diff --git a/HISTORY.md b/HISTORY.md index 608eee94..76ec9e9f 100644 --- a/HISTORY.md +++ b/HISTORY.md @@ -1,6 +1,6 @@ ## 0.15.4 -`value_and_gradient!!` and `value_gradient_and_hessian!!` now accept a `context=` keyword that overrides, for a single call, the `context` frozen at `prepare` — without re-preparing. The default (`context=nothing`) leaves the zero-allocation hot path unchanged; a `Tuple` of matching element types and shapes reuses the type-keyed cache. Two backends bake the context into their prepared state and throw an `ArgumentError` for a non-empty-input override (re-`prepare` instead): compiled-tape ReverseDiff (`AutoReverseDiff(; compile=true)`) and Mooncake's Hessian. Empty input runs no derivative machinery, so an override is always accepted there. +`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 is validated against the frozen context — matching element types and shapes, `ArgumentError` otherwise — 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. diff --git a/docs/src/evaluators.md b/docs/src/evaluators.md index 30d119fd..a70fa6f7 100644 --- a/docs/src/evaluators.md +++ b/docs/src/evaluators.md @@ -209,8 +209,8 @@ 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!!` or `value_gradient_and_hessian!!`, -without re-preparing: +`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)) @@ -219,13 +219,19 @@ value_and_gradient!!(prepared, x; context=(3.0, 0.0)) # overrides, this call on ``` The override must match the prepared context's element types and shapes (the -prepared cache is keyed on types), and it is per-call — it does not mutate the -frozen context. Two backends bake the context into their prepared state and -reject an override on a non-empty input with an `ArgumentError` (re-`prepare` +prepared cache is keyed on types); every override is validated against the +frozen context, and a 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 the tape, and Mooncake's `value_gradient_and_hessian!!`, -whose Hessian cache binds its target by object identity. Empty input runs no -derivative machinery, so an override is always accepted there. +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 diff --git a/ext/AbstractPPLDifferentiationInterfaceExt.jl b/ext/AbstractPPLDifferentiationInterfaceExt.jl index 35768f90..9686e962 100644 --- a/ext/AbstractPPLDifferentiationInterfaceExt.jl +++ b/ext/AbstractPPLDifferentiationInterfaceExt.jl @@ -178,16 +178,19 @@ end 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, _eval, ::Nothing @@ -242,23 +245,27 @@ 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 diff --git a/ext/AbstractPPLForwardDiffExt.jl b/ext/AbstractPPLForwardDiffExt.jl index 0e532064..a0a5399b 100644 --- a/ext/AbstractPPLForwardDiffExt.jl +++ b/ext/AbstractPPLForwardDiffExt.jl @@ -185,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 diff --git a/ext/AbstractPPLMooncakeExt.jl b/ext/AbstractPPLMooncakeExt.jl index dade72f0..f21d7e3b 100644 --- a/ext/AbstractPPLMooncakeExt.jl +++ b/ext/AbstractPPLMooncakeExt.jl @@ -283,7 +283,8 @@ end <:VectorEvaluator, <:Union{MooncakeCache{:scalar},MooncakeCache{:hessian}}, }, - ::AbstractVector{<:Real}, + ::AbstractVector{<:Real}; + context=nothing, ) return Evaluators._throw_jacobian_needs_vector() end @@ -292,20 +293,24 @@ 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)` 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 diff --git a/ext/AbstractPPLTestExt.jl b/ext/AbstractPPLTestExt.jl index 7923f3fa..a8227185 100644 --- a/ext/AbstractPPLTestExt.jl +++ b/ext/AbstractPPLTestExt.jl @@ -96,6 +96,25 @@ _context_problem(y::AbstractVector{<:Real}, offset) = -0.5 * (y[1] - offset)^2 # 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( @@ -317,6 +336,23 @@ function AbstractPPL.generate_testcases(::Val{:context_override}) 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 @@ -423,13 +459,18 @@ function _run( return nothing end -# Context frozen at `prepare` vs a per-call `context=` override, on both the -# gradient 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); each backend driver passes -# the pair that matches its caches. Empty input runs no derivative machinery, so -# an override is always accepted there regardless of the flags. +# 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/shape/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; @@ -440,6 +481,7 @@ function _run( check_dims::Bool=true, gradient_override::Symbol=:honor, hessian_override::Symbol=:honor, + jacobian_override::Symbol=:honor, kwargs..., ) ov = case.override @@ -506,6 +548,67 @@ function _run( 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/shape 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_shape = ntuple( + j -> j == i ? case.context[j][1:(end - 1)] : case.context[j], + length(case.context), + ) + @test_throws ArgumentError AbstractPPL.value_and_gradient!!( + prep, case.x; context=bad_shape + ) + 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` @@ -524,6 +627,15 @@ function _run( 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 + + # 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 diff --git a/src/evaluators/Evaluators.jl b/src/evaluators/Evaluators.jl index 2f71da5b..c5a91026 100644 --- a/src/evaluators/Evaluators.jl +++ b/src/evaluators/Evaluators.jl @@ -122,21 +122,31 @@ internal cache buffers of `prepared`. The returned gradient may alias 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. The override is per-call and does not mutate the frozen context. +keyed on types. Every override is validated against the frozen context; a +mismatch, or a non-`Tuple` override, throws an `ArgumentError`. The override is +per-call and does not mutate the frozen context. 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 a `context` override is always accepted there. +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 @@ -154,8 +164,9 @@ 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 an -override is always accepted there. +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 @@ -267,20 +278,77 @@ _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` calls `f` with the override directly (the per-call -# shape check having already run via `_check_ad_input` at the entry point). +# `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`). @inline _evaluate_with_context(e::VectorEvaluator, x, ::Nothing) = e(x) -@inline _evaluate_with_context(e::VectorEvaluator, x, context::Tuple) = e.f(x, context...) +@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 +# A non-`Tuple` override fails with the same `ArgumentError` as the AD paths. +@inline _evaluate_with_context(e::VectorEvaluator, x, context) = + _resolve_context(e, context) # 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 lives 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 -# and shapes, since the prepared cache is keyed on types. +# 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 and shapes, 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, +# and array elements are additionally compared by `size` (top-level elements +# only, mirroring the granularity of the backends' own cache validation). The +# `nothing` hot path is untouched. @inline _resolve_context(e::VectorEvaluator, ::Nothing) = e.context -@inline _resolve_context(::VectorEvaluator, context::Tuple) = context +@inline function _resolve_context(e::VectorEvaluator, context::Tuple) + typeof(context) === typeof(e.context) || + _throw_context_type_mismatch(e.context, context) + _check_context_sizes(e.context, context) + return context +end +function _resolve_context(::VectorEvaluator, context) + throw( + ArgumentError( + "A call-time `context` override must be a `Tuple`; got $(typeof(context))." + ), + ) +end + +function _throw_context_type_mismatch(frozen, override) + throw( + ArgumentError( + "Call-time `context` override does not match the context frozen at " * + "`prepare`: expected `$(typeof(frozen))`, got `$(typeof(override))`. " * + "The prepared caches are keyed on these types; re-`prepare` to change them.", + ), + ) +end + +@inline function _check_context_sizes(frozen::Tuple, override::Tuple) + ntuple(Val(length(frozen))) do i + a, b = frozen[i], override[i] + a isa AbstractArray && size(a) != size(b) && _throw_context_size_mismatch(i, a, b) + nothing + end + return nothing +end + +function _throw_context_size_mismatch(i, a, b) + throw( + ArgumentError( + "Call-time `context` override element $i has size $(size(b)); the " * + "context frozen at `prepare` has size $(size(a)). Re-`prepare` to " * + "change context shapes.", + ), + ) +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 diff --git a/test/ext/differentiationinterface/main.jl b/test/ext/differentiationinterface/main.jl index bb811809..d61dbf5a 100644 --- a/test/ext/differentiationinterface/main.jl +++ b/test/ext/differentiationinterface/main.jl @@ -101,9 +101,10 @@ quadratic(x::AbstractVector{<:Real}) = sum(xi -> xi^2, x) end # Non-compiled ReverseDiff threads context as call-time `Constant`s and - # honours an override on both entry points; compiled-tape ReverseDiff bakes - # context into its tape and rejects a non-empty override. Empty input is - # accepted on both (the runner checks it regardless of the flags). + # 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) @@ -114,6 +115,7 @@ quadratic(x::AbstractVector{<:Real}) = sum(xi -> xi^2, x) rtol=1e-6, gradient_override=:reject, hessian_override=:reject, + jacobian_override=:reject, ) end end diff --git a/test/ext/mooncake/main.jl b/test/ext/mooncake/main.jl index dc306297..43cb0629 100644 --- a/test/ext/mooncake/main.jl +++ b/test/ext/mooncake/main.jl @@ -186,7 +186,12 @@ end ) for case in generate_testcases(Val(:context_override)) run_testcase( - case; adtype=ad, check_dims=false, rtol=1.0e-6, hessian_override=:reject + case; + adtype=ad, + check_dims=false, + rtol=1.0e-6, + hessian_override=:reject, + jacobian_override=:prepare_rejects, ) end end