Call-time context overrides (#167); harden reused reverse-mode caches

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 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.
+
 ## 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:

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"

docs/src/evaluators.md

@@ -206,6 +206,33 @@ 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); 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 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

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