Add cache-reuse tests; validate NamedTuple shape on AD entry

- `:cache_reuse` conformance group in `AbstractPPLTestExt` drives
  `value_and_{gradient,jacobian}!!` three times per case against a single
  `prepared` evaluator to catch backend cache corruption between calls.
- `_assert_namedtuple_shape` always validates regardless of `CheckInput`;
  the `{false}` callable opts out by not calling here. Fixes a silent
  wrong-shape gradient when `prepare(...; check_dims=false)` was used with
  a NamedTuple input whose nested array sizes differ from the prototype.
  Mooncake regression test added.
- DI ext sub-environment now also loads `ReverseDiff` and exercises
  `AutoReverseDiff(compile=true)` against the conformance suite, covering
  the `_prepare_di(::AutoReverseDiff{true}, …)` compiled-tape path.
- Trimmed verbose `check_dims` clarifications in docstrings and
  `docs/src/evaluators.md` to one sentence each.

Co-Authored-By: Claude Opus 4.7 (1M context) <noreply@anthropic.com>

Author: yebai

docs/src/evaluators.md

@@ -138,6 +138,11 @@ library invokes the inner callable many times with same-length dual arrays
 derived from a single user-supplied `x`; re-validating on each invocation
 would be redundant work in the hot path.
 
+`check_dims=false` only disables the inner-evaluator hot path. Public AD
+entry points still validate input shape against the prepared evaluator —
+backend caches are tied to the prepared shape, so mismatches would otherwise
+silently produce wrong-shape derivatives.
+
 ## Without an AD backend
 
 The two-argument form `prepare(problem, x)` is available without any AD

ext/AbstractPPLTestExt.jl

@@ -232,4 +232,37 @@ function AbstractPPL.run_testcases(
     return nothing
 end
 
+# Drive `value_and_{gradient,jacobian}!!` twice with different inputs against
+# the same `prepared` evaluator to exercise cache reuse — catches backends
+# whose cache state is corrupted by a prior call.
+function AbstractPPL.run_testcases(
+    ::Val{:cache_reuse}, prepare_fn=AbstractPPL.prepare; adtype, atol=0, rtol=1e-10
+)
+    @testset "scalar output, repeated calls" begin
+        prepared = prepare_fn(adtype, QuadraticProblem(), zeros(3))
+        for (x, value, gradient) in (
+            ([1.0, 2.0, 3.0], 14.0, [2.0, 4.0, 6.0]),
+            ([4.0, 5.0, 6.0], 77.0, [8.0, 10.0, 12.0]),
+            ([0.5, -1.0, 2.0], 5.25, [1.0, -2.0, 4.0]),
+        )
+            val, grad = AbstractPPL.value_and_gradient!!(prepared, x)
+            @test val ≈ value atol = atol rtol = rtol
+            @test grad ≈ gradient atol = atol rtol = rtol
+        end
+    end
+    @testset "vector output, repeated calls" begin
+        prepared = prepare_fn(adtype, VectorValuedProblem(), zeros(3))
+        for (x, value, jacobian) in (
+            ([2.0, 3.0, 4.0], [6.0, 7.0], [3.0 2.0 0.0; 0.0 1.0 1.0]),
+            ([5.0, 1.0, 7.0], [5.0, 8.0], [1.0 5.0 0.0; 0.0 1.0 1.0]),
+            ([0.0, 4.0, -2.0], [0.0, 2.0], [4.0 0.0 0.0; 0.0 1.0 1.0]),
+        )
+            val, jac = AbstractPPL.value_and_jacobian!!(prepared, x)
+            @test val ≈ value atol = atol rtol = rtol
+            @test jac ≈ jacobian atol = atol rtol = rtol
+        end
+    end
+    return nothing
+end
+
 end # module

src/evaluators/Evaluators.jl

@@ -54,7 +54,9 @@ prepares gradient or jacobian machinery for vector inputs.
 `check_dims` (default `true`) controls whether the returned evaluator validates
 the input shape on each call. Pass `check_dims=false` to skip the per-call
 check, e.g. inside an AD backend's hot path where the input shape is already
-guaranteed.
+guaranteed. AD entry points (`value_and_gradient!!` / `value_and_jacobian!!`)
+still validate input shape against the prepared evaluator regardless — backend
+caches are tied to the prepared shape.
 
 The three-argument AD-aware form may invoke `problem` once during preparation
 to detect output arity (scalar vs vector) and select gradient or jacobian
@@ -145,6 +147,10 @@ prototype's `typeof` is captured at preparation time using the original element
 types, so a `CheckInput=true` evaluator will reject inputs whose leaves are
 dual/shadow numbers (or any other widened element type) even when the structure
 is otherwise correct.
+
+`CheckInput=false` only disables the inner-evaluator check; public AD entry
+points (`value_and_gradient!!`) always validate via `_assert_namedtuple_shape`
+to keep cache reuse safe.
 """
 struct NamedTupleEvaluator{CheckInput,F,P<:NamedTuple}
     f::F
@@ -208,10 +214,12 @@ end
 Throw `ArgumentError` unless `values` has the same type as the prototype captured
 during preparation, including matching `size` for any nested `AbstractArray`
 leaves. Also throws if the prototype contains a leaf type outside the supported
-set (`Real`, `Complex`, `AbstractArray`, `Tuple`, `NamedTuple`). No-op when `e`
-was constructed with `CheckInput=false`.
+set (`Real`, `Complex`, `AbstractArray`, `Tuple`, `NamedTuple`).
+
+Always validates regardless of the evaluator's `CheckInput` parameter; the
+`{false}` callable opts out by not calling here.
 """
-function _assert_namedtuple_shape(e::NamedTupleEvaluator{true}, values)
+function _assert_namedtuple_shape(e::NamedTupleEvaluator, values)
     typeof(values) === typeof(e.inputspec) || throw(
         ArgumentError(
             "Expected the same NamedTuple structure that was used to prepare this evaluator.",
@@ -224,7 +232,6 @@ function _assert_namedtuple_shape(e::NamedTupleEvaluator{true}, values)
     )
     return nothing
 end
-_assert_namedtuple_shape(::NamedTupleEvaluator{false}, _) = nothing
 
 # Classify the output of a probe `evaluator(x)` call into the two arities the
 # AD interface supports — `:scalar` routes to gradient prep, `:vector` to

test/ext/differentiationinterface/Project.toml

@@ -4,10 +4,12 @@ ADTypes = "47edcb42-4c32-4615-8424-f2b9edc5f35b"
 DifferentiationInterface = "a0c0ee7d-e4b9-4e03-894e-1c5f64a51d63"
 ForwardDiff = "f6369f11-7733-5829-9624-2563aa707210"
 Pkg = "44cfe95a-1eb2-52ea-b672-e2afdf69b78f"
+ReverseDiff = "37e2e3b7-166d-5795-8a7a-e32c996b4267"
 Test = "8dfed614-e22c-5e08-85e1-65c5234f0b40"
 
 [compat]
 ADTypes = "1"
 DifferentiationInterface = "0.6, 0.7"
 ForwardDiff = "1"
+ReverseDiff = "1"
 julia = "1.10"

test/ext/differentiationinterface/main.jl

@@ -4,12 +4,26 @@ Pkg.develop(; path=joinpath(@__DIR__, "..", "..", ".."))
 Pkg.instantiate()
 
 using AbstractPPL: run_testcases
-using ADTypes: AutoForwardDiff
+using ADTypes: AutoForwardDiff, AutoReverseDiff
 using DifferentiationInterface: DifferentiationInterface as DI
 using ForwardDiff
+using ReverseDiff
 using Test
 
 @testset "AbstractPPLDifferentiationInterfaceExt" begin
-    run_testcases(Val(:vector); adtype=AutoForwardDiff(), atol=1e-6, rtol=1e-6)
-    run_testcases(Val(:edge); adtype=AutoForwardDiff())
+    @testset "ForwardDiff" begin
+        run_testcases(Val(:vector); adtype=AutoForwardDiff(), atol=1e-6, rtol=1e-6)
+        run_testcases(Val(:cache_reuse); adtype=AutoForwardDiff(), atol=1e-6, rtol=1e-6)
+        run_testcases(Val(:edge); adtype=AutoForwardDiff())
+    end
+
+    # Compiled-tape ReverseDiff goes through the `_prepare_di(::AutoReverseDiff{true}, …)`
+    # specialisation that closes the evaluator into a `Base.Fix2` target — the
+    # `:cache_reuse` group exercises that path across multiple inputs.
+    @testset "ReverseDiff (compiled tape)" begin
+        adtype = AutoReverseDiff(; compile=true)
+        run_testcases(Val(:vector); adtype=adtype, atol=1e-6, rtol=1e-6)
+        run_testcases(Val(:cache_reuse); adtype=adtype, atol=1e-6, rtol=1e-6)
+        run_testcases(Val(:edge); adtype=adtype)
+    end
 end

test/ext/mooncake/main.jl

@@ -16,7 +16,21 @@ using Test
         @testset "$label" begin
             run_testcases(Val(:vector); adtype=adtype, atol=1e-6, rtol=1e-6)
             run_testcases(Val(:namedtuple); adtype=adtype, atol=1e-6, rtol=1e-6)
+            run_testcases(Val(:cache_reuse); adtype=adtype, atol=1e-6, rtol=1e-6)
             run_testcases(Val(:edge); adtype=adtype)
         end
     end
+
+    # Regression: `check_dims=false` skips the per-call shape check on the AD
+    # hot path, but the public `value_and_gradient!!` must still validate cache
+    # compatibility — Mooncake caches are tied to the prototype's nested array
+    # sizes, so a mismatched leaf would silently return a wrong-shape gradient.
+    @testset "NamedTuple cache validates regardless of check_dims" begin
+        f = vs -> vs.x^2 + sum(abs2, vs.y)
+        proto = (x=0.0, y=zeros(2))
+        prepared = prepare(AutoMooncake(), f, proto; check_dims=false)
+        @test_throws r"Nested array shape" AbstractPPL.value_and_gradient!!(
+            prepared, (x=3.0, y=[1.0])
+        )
+    end
 end